jax-ml/jax · error · TypeError

Number of tensordot axes (axes {}) exceeds input ranks ({} a

Error message

Number of tensordot axes (axes {}) exceeds input ranks ({} and {})

What it means

jax.numpy.tensordot raises this TypeError when the integer axes argument is larger than the number of dimensions of either input array. The contraction rank cannot exceed min(a.ndim, b.ndim), so the operation is mathematically ill-formed and JAX rejects it before dispatching to lax.dot_general.

Source

Thrown at jax/_src/numpy/tensor_contractions.py:566

    >>> jnp.outer(x1, x2)
    Array([[1, 2, 3],
           [2, 4, 6]], dtype=int32)
  """
  a, b = util.ensure_arraylike("tensordot", a, b)
  a_ndim = np.ndim(a)
  b_ndim = np.ndim(b)

  if preferred_element_type is None:
    preferred_element_type, output_weak_type = dtypes.result_type(a, b, return_weak_type_flag=True)
  else:
    preferred_element_type = dtypes.check_and_canonicalize_user_dtype(
        preferred_element_type, "tensordot")
    output_weak_type = False

  if type(axes) is int:
    if axes > min(a_ndim, b_ndim):
      msg = "Number of tensordot axes (axes {}) exceeds input ranks ({} and {})"
      raise TypeError(msg.format(axes, a.shape, b.shape))
    contracting_dims = tuple(range(a_ndim - axes, a_ndim)), tuple(range(axes))
  elif isinstance(axes, (tuple, list)) and len(axes) == 2:
    ax1, ax2 = axes
    if isinstance(ax1, int) and isinstance(ax2, int):
      contracting_dims = ((canonicalize_axis(ax1, a_ndim),),
                          (canonicalize_axis(ax2, b_ndim),))
    elif isinstance(ax1, (tuple, list)) and isinstance(ax2, (tuple, list)):
      if len(ax1) != len(ax2):
        msg = "tensordot requires axes lists to have equal length, got {} and {}."
        raise TypeError(msg.format(ax1, ax2))
      contracting_dims = (tuple(canonicalize_axis(i, a_ndim) for i in ax1),
                          tuple(canonicalize_axis(i, b_ndim) for i in ax2))
    else:
      msg = ("tensordot requires both axes lists to be either ints, tuples or "
             "lists, got {} and {}")
      raise TypeError(msg.format(ax1, ax2))
  else:
    msg = ("tensordot axes argument must be an int, a pair of ints, or a pair "

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Check a.ndim and b.ndim before calling and pass axes <= min(a.ndim, b.ndim)
2. Fix the inputs so both arrays have at least the expected rank (e.g. add batch dims or reshape)
3. Pass explicit axis pairs instead of an integer count, e.g. axes=([1],[0])

Example fix

// before
jnp.tensordot(a, b, axes=3)  # a.ndim == 2
// after
jnp.tensordot(a, b, axes=2)  # or axes=([1],[0])

assert isinstance(axes, int) and axes <= min(a.ndim, b.ndim), f"axes={axes} exceeds ranks {a.ndim},{b.ndim}"

When it happens

Trigger: Calling jnp.tensordot(a, b, axes=N) with N > min(a.ndim, b.ndim), e.g. jnp.tensordot(jnp.zeros((2,3)), jnp.zeros((3,4)), axes=3).

Common situations: Mismatch between the intended matrix/tensor ranks (e.g. passing a scalar or 1D vector where a 2D matrix was expected), or copying numpy code that assumed higher-rank batches.

Related errors

• multi_dot: last dimension of each array must match first dim
• expected x and y to have same length
• Weights shape must match 'a' shape when axis is None.
• Weights shape {weights.shape} must match reduction axes {tup
• tensordot requires axes lists to have equal length, got {} a

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