jax-ml/jax · error · NotImplementedError

Bitcast ref with dynamic size is not supported.

Error message

Bitcast ref with dynamic size is not supported.

What it means

TransformedRef.bitcast raises NotImplementedError when the ref has a dynamic size (shape containing tracers/int32 symbolic dims), because bitcasting requires static sizes to compute the new shape.

Source

Thrown at jax/_src/state/types.py:354

    return self.type.shape

  @property
  def dtype(self):
    if not hasattr(self.type, "dtype"):
      raise AttributeError(f"{self!r} has no `dtype`.") from None
    return self.type.dtype

  ndim = property(lambda self: len(self.shape))
  size = property(lambda self: math.prod(self.shape))
  T = property(lambda self: self.transpose(tuple(reversed(range(self.ndim)))))

  @property
  def at(self) -> RefIndexer:
    return RefIndexer(self)

  def bitcast(self, dtype):
    if self.is_dynamic_size:
      raise NotImplementedError(
          "Bitcast ref with dynamic size is not supported."
      )
    dtype = dtypes.dtype(dtype)
    if self.multiref:
      return TransformedRef(self, (BitcastTransform(dtype),))
    return TransformedRef(self.ref, (*self.transforms, BitcastTransform(dtype)))

  def reshape(self, *shape):
    if self.is_dynamic_size:
      raise NotImplementedError(
          "Reshape ref with dynamic size is not supported."
      )
    if len(shape) == 1 and isinstance(shape[0], tuple):
      shape = shape[0]
    input_shape = tuple(operator.index(s) for s in self.shape)
    shape = _canonicalize_reshape(input_shape, shape)
    if self.multiref:
      return TransformedRef(self, (ReshapeTransform(shape),))

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Hoist the bitcast outside the dynamic-shape region
2. Use static shapes for that portion of the computation
3. Reinterpret bytes manually via get/put with converted values

if ref.is_dynamic_size:
    raise RuntimeError("bitcast unsupported here; hoist out of dynamic region")

When it happens

Trigger: Calling .bitcast(dtype) on a ref created inside a transform with dynamic shapes (e.g. inside remat/pmap or with polymorphic shapes).

Common situations: Using bitcast inside dynamically shaped computations; migration from static-shape code to dynamic shapes.

Related errors

• associative scan over axis of non-constant size: {}. You may
• can only convert to an extended dtype on an array type,but g
• Attempting to convert array of shape {operand.shape} from {o
• lax.bitcast_convert_type does not support bool or complex va
• bitcast_convert_type with different bitwidths not supported

AI-assisted analysis of jax-ml/jax@1e1c6a8fc0 (2026-08-27). Data as JSON: /api/errors/e6e8b81f17b3cadc. Report an issue: GitHub.