jax-ml/jax · error · ValueError

function is not returning an array of the correct shape

Error message

function is not returning an array of the correct shape

What it means

Raised by jnp.apply_along_axis (the apply-like helper that reinflates func output): the user function must return either the same number of dims as the input slice (ndim preserved) or exactly one fewer (it is re-expanded along axis). Any other ndim triggers this error.

Source

Thrown at jax/_src/numpy/lax_numpy.py:7960

           [4]], dtype=int32)
    >>> jnp.min(x, [1], keepdims=True)
    Array([[1],
           [4]], dtype=int32)

    >>> jnp.apply_over_axes(jnp.prod, x, [0, 1])
    Array([[720]], dtype=int32)
    >>> jnp.prod(x, [0, 1], keepdims=True)
    Array([[720]], dtype=int32)
  """
  a_arr = util.ensure_arraylike("apply_over_axes", a)
  for axis in axes:
    b = func(a_arr, axis)
    if b.ndim == a_arr.ndim:
      a_arr = b
    elif b.ndim == a_arr.ndim - 1:
      a_arr = expand_dims(b, axis)
    else:
      raise ValueError("function is not returning an array of the correct shape")
  return a_arr


@export
@api.jit(static_argnames=('axisa', 'axisb', 'axisc', 'axis'))
def cross(a, b, axisa: int = -1, axisb: int = -1, axisc: int = -1,
          axis: int | None = None):
  r"""Compute the (batched) cross product of two arrays.

  JAX implementation of :func:`numpy.cross`.

  This computes the 2-dimensional or 3-dimensional cross product,

  .. math::

     c = a \times b

  In 3 dimensions, ``c`` is a length-3 array. In 2 dimensions, ``c`` is

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Make func return exactly a 1-D array per 1-D slice (or matching-ndim output)
2. Wrap scalar returns with at least 1-D: ensure output ndim == input slice ndim or ndim-1 consistently
3. Use jax.vmap instead of apply_along_axis, which handles arbitrary output shapes

Example fix

// before
jnp.apply_along_axis(lambda row: row[:, None] @ row[None, :], 1, a)  # returns 2-D
// after
jax.vmap(lambda row: row[:, None] @ row[None, :], in_axes=0)(a)

out = func(a[0])  # probe one slice
assert jnp.asarray(out).ndim in (a[0].ndim, a[0].ndim - 1)

def returns_compatible_shape(func, sample_slice):
    b = jnp.asarray(func(sample_slice))
    return b.ndim in (sample_slice.ndim, sample_slice.ndim - 1)

When it happens

Trigger: Passing a function to jnp.apply_along_axis that returns a 2-D result from 1-D slices, or a scalar from 2-D slices where the plumbing expects ndim or ndim-1; returning a list/tuple whose asarray ndim differs unexpectedly.

Common situations: The mapped function's return shape changes with data (e.g. returns [] for some rows); applying a function that returns multiple values packed in extra dims; refactoring a function so it now returns a tuple wrapped array.

Related errors

• unexpected JAX type (e.g. shape/dtype) for argument to VJP f
• cotangent type does not match function output, expected {out
• Mismatched number of outputs from callback. Expected: {}, Ac
• Incorrect output shape for return value #{i}: Expected: {out
• Length of sharding.spec ({len(out_s.spec)}) must be equal to

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