jax-ml/jax · error · TypeError

tensordot requires axes lists to have equal length, got {} a

Error message

tensordot requires axes lists to have equal length, got {} and {}.

What it means

When tensordot's axes argument is a pair of sequences, both sequences must have the same length because each entry of ax1 contracts with the corresponding entry of ax2. JAX raises this TypeError when len(ax1) != len(ax2), since the contraction is undefined.

Source

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

  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 "
           "of lists/tuples of ints.")
    raise TypeError(msg)
  result = lax.dot_general(
      a, b, (contracting_dims, ((), ())), precision=precision,
      preferred_element_type=preferred_element_type,
      out_sharding=out_sharding)
  return lax._convert_element_type(result, preferred_element_type, output_weak_type)

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Make both axis lists the same length
2. Verify each axis index is within range of the respective array's ndim

Example fix

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

if isinstance(axes, tuple) and len(axes) == 2 and isinstance(axes[0], (list, tuple)):
    assert len(axes[0]) == len(axes[1]), 'axis lists must match'

When it happens

Trigger: jnp.tensordot(a, b, axes=([0,1],[2])) where the two axis lists differ in length.

Common situations: Typos or off-by-one in hand-written axis lists; refactoring code where an axis was added to one side only.

Related errors

• Number of tensordot axes (axes {}) exceeds input ranks ({} a
• tensordot axes argument must be an int, a pair of ints, or a
• len() of unsized object
• numpy masked arrays are not supported as direct inputs to JA
• Unsupported scalar attribute type: {type(val)}

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