jax-ml/jax · error · ValueError

Dimension must be either 2 or 3 for cross product

Error message

Dimension must be either 2 or 3 for cross product

What it means

Raised by jnp.cross when, after moving the specified axes to the last position, either input's last dimension is not 2 or 3 — cross products are only defined for 2-D and 3-D vectors in JAX.

Source

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

    >>> a = jnp.array([[1, 2, 3],
    ...                [3, 4, 3]])
    >>> b = jnp.array([[2, 3, 2],
    ...                [4, 5, 6]])
    >>> jnp.cross(a, b)
    Array([[-5,  4, -1],
           [ 9, -6, -1]], dtype=int32)
  """
  util.check_arraylike("cross", a, b)
  if axis is not None:
    axisa = axis
    axisb = axis
    axisc = axis
  a = moveaxis(a, axisa, -1)
  b = moveaxis(b, axisb, -1)

  if a.shape[-1] not in (2, 3) or b.shape[-1] not in (2, 3):
    raise ValueError("Dimension must be either 2 or 3 for cross product")

  if a.shape[-1] == 2 or b.shape[-1] == 2:
    deprecations.warn(
        "jax-numpy-cross-2d-input",
        "Support for 2-dimensional vectors in jnp.cross is deprecated and "
        "will be removed in JAX 0.12.0. Use arrays of 3-dimensional "
        "vectors instead.",
        stacklevel=2,
    )

  if a.shape[-1] == 2 and b.shape[-1] == 2:
    return a[..., 0] * b[..., 1] - a[..., 1] * b[..., 0]

  a0 = a[..., 0]
  a1 = a[..., 1]
  a2 = a[..., 2] if a.shape[-1] == 3 else array_creation.zeros_like(a0)
  b0 = b[..., 0]
  b1 = b[..., 1]

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

  1. Fix axisa/axisb (or axis) so the vector dimension (2 or 3) is the axis used
  2. Slice/pad inputs to exactly 2 or 3 components (e.g. pad 2-D to 3-D with zeros)
  3. Ensure inputs are at least 1-D with shape[-1] in (2, 3) before calling

Example fix

// before
jnp.cross(a, b, axis=0)  # data stacked as (3, N), axis=-1 intended
// after
jnp.cross(a, b, axis=0)  # verify shape[0] == 3; or transpose: jnp.cross(a.T, b.T)
Defensive patterns

Strategy: validation

Validate before calling

assert a.shape[axisa] in (2, 3) and b.shape[axisb] in (2, 3)

Type guard

def crossable(a, b, axisa=-1, axisb=-1):
    return jnp.moveaxis(a, axisa, -1).shape[-1] in (2, 3) and jnp.moveaxis(b, axisb, -1).shape[-1] in (2, 3)

Prevention

When it happens

Trigger: jnp.cross(a, b) where a.shape[-1] is 1, 4, or the arrays are scalars/empty on the last axis; wrong axisa/axisb so a non-vector axis ends up last; passing stacked matrices whose leading dim is misinterpreted.

Common situations: Using axis=0 with column-stacked data so shape[-1] != 2/3; passing 4-D homogeneous coordinates; porting code that used scipy cross-free helpers.

Related errors


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