Method bind

jax/_src/core.py:660–707

(self, *args, **params)

    return f'{self.name}'

  def bind(self, *args, **params):
    canonical_args = []
    avals = []
    for i, arg in enumerate(args):
      try:
        c_arg = dtypes.canonicalize_value(arg)
        aval = typeof(c_arg)
      except TypeError as e:
        raise TypeError(
          f"Error interpreting argument to {self} as a JAX value."
          f" The problematic value is of type {type(arg)} and was passed to"
          f" {self} at position {i}.\n"
        ) from e
      if (not self.skip_canonicalization and isinstance(aval, ShapedArray)
          and not aval.sharding.mesh.empty):
        cur_mesh = mesh_lib.get_abstract_mesh()
        if cur_mesh != aval.sharding.mesh:
          # TODO(yashkatariya): Casting to Explicit is not yet allowed. Maybe we
          # need cast_and_slice_p for it since shape might change?
          # Atleast 1 mesh axis should be Manual and all other axes should be
          # Manual or Auto to allow casting.
          if cur_mesh._any_axis_manual and cur_mesh._are_all_axes_auto_or_manual:
            if aval.sharding.mesh.are_all_axes_auto:
              from jax._src.pjit import reshard  # pyrefly: ignore[missing-import]
              c_arg = reshard(c_arg, NamedSharding(cur_mesh, P(*[None] * aval.ndim)))
              aval = typeof(c_arg)
            elif aval.sharding.mesh._any_axis_explicit:
              raise NotImplementedError(
                  "Closing over inputs to shard_map where the input is sharded "
                  "on `Explicit` axes is not implemented. As a workaround, "
                  "please pass those inputs as an argument to shard_map. Got "
                  f"input with shape {aval.str_short(True, True)}")
      if isinstance(c_arg, Tracer) and not c_arg._trace.is_valid():
        raise escaped_tracer_error(c_arg)
      canonical_args.append(c_arg)
      avals.append(aval)

    args = canonical_args

    # This is equivalent to "with take_current_trace()", but the bind() code
    # is called frequently and it's slightly faster to avoid using a context
    # manager object.
    prev_trace = trace_ctx.trace
    trace_ctx.set_trace(eval_trace)
    try:
      return self.bind_with_trace(prev_trace, args, avals, params)
    finally:
      trace_ctx.set_trace(prev_trace)

  def bind_with_trace(self, trace, args, avals, params, /):
    if self.is_high(*avals, **params) and trace.requires_low: