jax-ml/jax · error · NotImplementedError

Reshape ref with dynamic size is not supported.

Error message

Reshape ref with dynamic size is not supported.

What it means

TransformedRef.reshape raises NotImplementedError when the ref has a dynamic size, since canonicalizing the reshape requires concrete input and output shapes.

Source

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

  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),))
    return TransformedRef(self.ref, (*self.transforms, ReshapeTransform(shape)))

  def transpose(self, permutation: Sequence[int]):
    if self.multiref:
      raise NotImplementedError("Transpose with multiref is not supported.")
    transposer = TransposeTransform(tuple(permutation))
    if self.multiref:
      return TransformedRef(self, (transposer,))
    return TransformedRef(self.ref, (*self.transforms, transposer))

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Move the reshape outside the dynamic region
2. Use static shapes for the buffer
3. Allocate the target shape directly instead of reshaping

if ref.is_dynamic_size:
    raise RuntimeError("reshape unsupported here; use static shapes")

When it happens

Trigger: Calling .reshape(...) on a ref with symbolic/dynamic dimensions.

Common situations: Reshaping buffers inside shape-polymorphic code paths.

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
• reshape new_sizes must all be positive, got {}.
• np.reshape order=A is not implemented.
• Unexpected value for 'order' argument: {order}.

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