{"record":{"id":"00320029a8a6a3f2","repo":"jax-ml/jax","slug":"function-is-not-returning-an-array-of-the-correct","errorCode":null,"errorMessage":"function is not returning an array of the correct shape","messagePattern":"function is not returning an array of the correct shape","errorType":"exception","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/numpy/lax_numpy.py","lineNumber":7960,"sourceCode":"           [4]], dtype=int32)\n    >>> jnp.min(x, [1], keepdims=True)\n    Array([[1],\n           [4]], dtype=int32)\n\n    >>> jnp.apply_over_axes(jnp.prod, x, [0, 1])\n    Array([[720]], dtype=int32)\n    >>> jnp.prod(x, [0, 1], keepdims=True)\n    Array([[720]], dtype=int32)\n  \"\"\"\n  a_arr = util.ensure_arraylike(\"apply_over_axes\", a)\n  for axis in axes:\n    b = func(a_arr, axis)\n    if b.ndim == a_arr.ndim:\n      a_arr = b\n    elif b.ndim == a_arr.ndim - 1:\n      a_arr = expand_dims(b, axis)\n    else:\n      raise ValueError(\"function is not returning an array of the correct shape\")\n  return a_arr\n\n\n@export\n@api.jit(static_argnames=('axisa', 'axisb', 'axisc', 'axis'))\ndef cross(a, b, axisa: int = -1, axisb: int = -1, axisc: int = -1,\n          axis: int | None = None):\n  r\"\"\"Compute the (batched) cross product of two arrays.\n\n  JAX implementation of :func:`numpy.cross`.\n\n  This computes the 2-dimensional or 3-dimensional cross product,\n\n  .. math::\n\n     c = a \\times b\n\n  In 3 dimensions, ``c`` is a length-3 array. In 2 dimensions, ``c`` is","sourceCodeStart":7942,"sourceCodeEnd":7978,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/numpy/lax_numpy.py#L7942-L7978","documentation":"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.","triggerScenarios":"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.","commonSituations":"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.","solutions":["Make func return exactly a 1-D array per 1-D slice (or matching-ndim output)","Wrap scalar returns with at least 1-D: ensure output ndim == input slice ndim or ndim-1 consistently","Use jax.vmap instead of apply_along_axis, which handles arbitrary output shapes"],"exampleFix":"// before\njnp.apply_along_axis(lambda row: row[:, None] @ row[None, :], 1, a)  # returns 2-D\n// after\njax.vmap(lambda row: row[:, None] @ row[None, :], in_axes=0)(a)\n","handlingStrategy":"validation","validationCode":"out = func(a[0])  # probe one slice\nassert jnp.asarray(out).ndim in (a[0].ndim, a[0].ndim - 1)","typeGuard":"def returns_compatible_shape(func, sample_slice):\n    b = jnp.asarray(func(sample_slice))\n    return b.ndim in (sample_slice.ndim, sample_slice.ndim - 1)","tryCatchPattern":null,"preventionTips":["Make the applied function return fixed-shape 1-D arrays","Prefer jax.vmap for arbitrary output shapes","Test the function on one slice before applying"],"tags":["jax","apply-along-axis","shape-mismatch"],"backgroundTag":"callback-return-shape-mismatch","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}