jax-ml/jax · error · TypeError

tensordot axes argument must be an int, a pair of ints, or a

Error message

tensordot axes argument must be an int, a pair of ints, or a pair of lists/tuples of ints.

What it means

tensordot accepts only three axes forms: a single int, a pair of ints, or a pair of lists/tuples of ints. Anything else (a bare list, a string, a 3-element tuple, etc.) hits the final else branch and raises this TypeError.

Source

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

  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)



@export
@api.jit(static_argnames=('precision', 'preferred_element_type'), inline=True)
def inner(
    a: ArrayLike, b: ArrayLike, *, precision: lax.PrecisionLike = None,
    preferred_element_type: DTypeLike | None = None,
) -> Array:
  """Compute the inner product of two arrays.

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

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Pass an integer N, or a (int,int) pair, or a (list,list) pair
2. Wrap a single axis list as ([...],[...])

Example fix

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

ok = type(axes) is int or (isinstance(axes,(tuple,list)) and len(axes)==2)
assert ok, 'axes must be int, (int,int), or (list,list)'

When it happens

Trigger: jnp.tensordot(a, b, axes=[1,2]) (a single list, not a pair), or axes=((1,2),(0,1),(3,4)), or a non-integer value.

Common situations: Confusion with numpy semantics or other libraries; passing a single sequence where a pair-of-sequences is required; passing None.

Related errors

• method can only be 'linear', 'lower', 'higher', 'midpoint',
• Number of tensordot axes (axes {}) exceeds input ranks ({} a
• tensordot requires axes lists to have equal length, got {} a
• len() of unsized object
• Unknown resize method "{s}"

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