jax-ml/jax · error · ValueError

Array type must have a `shape` attribute, but got {type(arra

Error message

Array type must have a `shape` attribute, but got {type(array_aval)}

What it means

BlockSpec.to_block_mapping needs the array's shape to build default index maps and block mappings. If block_shape is given but the input aval has no shape attribute (not an array aval — e.g. a token or a custom abstract value), it raises this error.

Source

Thrown at jax/_src/pallas/core.py:592

      #     traceback_util.api_boundary(self.index_map, repro_user_func=True))
      self.index_map = _IndexMapFunc(self.index_map)

  def to_block_mapping(
      self,
      origin: OriginStr,
      array_aval: jax_core.ShapedArray,
      *,
      # Inputs for the index_map
      index_map_avals: Sequence[jax_core.AbstractValue],
      index_map_tree: tree_util.PyTreeDef,
      grid: GridMappingGrid,
      vmapped_dims: tuple[int, ...],
      debug: bool = False,
      allow_captured_consts: bool = False,
  ) -> BlockMapping:
    if self.block_shape is not None:
      if not hasattr(array_aval, "shape"):
        raise ValueError(
            "Array type must have a `shape` attribute, but got"
            f" {type(array_aval)}"
        )
    if self.index_map is None:
      index_map_func = default_index_map(len(array_aval.shape))
      index_map_dbg = api_util.debug_info("pallas_call index_map",
                                          default_index_map, (),{}
                                          )._replace(arg_names=("",) * len(index_map_avals))
      api_util.save_wrapped_fun_debug_info(index_map_func, index_map_dbg)
    else:
      index_map_func = self.index_map
    if self.block_shape is None:
      block_shape = _canonicalize_block_shape(array_aval.shape)
    else:
      block_shape = _canonicalize_block_shape(self.block_shape)
      if len(array_aval.shape) != len(block_shape):
        raise ValueError(
            f"Block shape for {origin} (= {block_shape}) "

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Verify every operand/output paired with a BlockSpec is a real array (jnp.asarray-able)
2. Check argument order of pallas_call(f, out_shape, grid, in_specs, out_specs)
3. Pass block_shape=None for non-array operands if that is intended

assert hasattr(array_aval, 'shape'), f'operand aval {array_aval} is not an array'
# only attach non-None block_shape BlockSpecs to array operands

def is_array_aval(aval):
    return hasattr(aval, 'shape') and hasattr(aval, 'dtype')

When it happens

Trigger: Calling pallas_call with a BlockSpec having a non-None block_shape on an operand whose aval lacks .shape — e.g. donating a token, passing an odd object, or misordering out_shape and BlockSpec arguments.

Common situations: Mixing up positional arguments of pallas_call; using BlockSpecs with non-array outputs; internal state-discharge paths after API changes.

Related errors

• MemoryRef type must be a ShapedArray, got {type(self.inner_a
• Unsupported block dimension type: {type(dim)}. Allowed types
• Unsupported block shape type: {type(dim)}
• Block shape for {origin} (= {block_shape}) must have the sam
• Index map function {debug_info.func_src_info} for {origin} m

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