jax-ml/jax · error · TypeError

tensordot requires both axes lists to be either ints, tuples

Error message

tensordot requires both axes lists to be either ints, tuples or lists, got {} and {}

What it means

When axes is a pair, each element must be either an int or a tuple/list of ints. JAX raises this TypeError when one element is a mixed or unsupported type (e.g. one int and one list, or a string).

Source

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

    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)



@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:

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Normalize both elements to the same type (both ints or both lists/tuples of ints)
2. Convert numpy ints or 0-d arrays to Python int before passing

Example fix

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

ax1, ax2 = axes
assert (isinstance(ax1, int) and isinstance(ax2, int)) or (isinstance(ax1,(list,tuple)) and isinstance(ax2,(list,tuple)))

def valid_axes_pair(ax1, ax2):
    return (type(ax1) is int and type(ax2) is int) or (isinstance(ax1,(list,tuple)) and isinstance(ax2,(list,tuple)))

When it happens

Trigger: jnp.tensordot(a, b, axes=(0, [1,2])) or axes=('x','y') — one side is an int while the other is a sequence, or a non-int type.

Common situations: Dynamically constructed axes arguments where one branch yields an int and the other a list; passing numpy arrays or strings as axes.

Related errors

• {name} was requested to map a value of non-array type {core.
• primal and tangent arguments to jax.jvp must be tuples or li
• {prim_name} takes a scalar pred as argument, got {pred}
• `compute_on`'s compute_type argument must be a string.
• Cannot interpret value of type {typ} as an abstract array; i

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