jax-ml/jax · error · NotImplementedError

reshape not supported yet: {aval_in}, {aval_out}

Error message

reshape not supported yet: {aval_in}, {aval_out}

What it means

Error "reshape not supported yet: {aval_in}, {aval_out}" thrown in jax-ml/jax.

Source

Thrown at jax/_src/pallas/fuser/block_spec.py:2467

          'reshape with non-matching block size on lanes not supported yet:'
          f' {block_shape}'
      )
    new_block_shape = (*block_shape[:-2], total_block_size)

    def new_block_index_transform(*idxs):
      *idx, second_to_last, last = block_transform.block_index_transform(*idxs)
      # last should always be 0
      if not isinstance(last, int) and last != 0:
        raise NotImplementedError(
            'Must select entire block on last dimension for reshape'
        )
      return *idx, second_to_last

    return [block_transform.replace(
                block_shape=new_block_shape,
                block_index_transform=new_block_index_transform,)]

  raise NotImplementedError(f'reshape not supported yet: {aval_in}, {aval_out}')


@register_eval_rule(lax.reshape_p)
def _reshape_eval_rule(
    eval_ctx: KernelEvalContext, x, *, dimensions, new_sizes, sharding
):
  del sharding, dimensions, new_sizes
  out_shape_nones = tuple(
      _block_size(s) for s in eval_ctx.out_block_specs[0].block_shape
  )
  out_shape = tuple(s for s in out_shape_nones if s is not None)
  # Because we have restricted the pull block spec rule, we can just apply a
  # basic reshape here.
  x = x.reshape(out_shape)
  return x


@register_pull_block_spec_rule(lax.reduce_sum_p)

View on GitHub (pinned to 1e1c6a8fc0)

When it happens

Trigger: Thrown at jax/_src/pallas/fuser/block_spec.py:2467 when the library encounters an invalid state.

Common situations: See trigger scenarios.

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