{"record":{"id":"3f51351b6efb37af","repo":"jax-ml/jax","slug":"tensordot-requires-both-axes-lists-to-be-either-in","errorCode":null,"errorMessage":"tensordot requires both axes lists to be either ints, tuples or lists, got {} and {}","messagePattern":"tensordot requires both axes lists to be either ints, tuples or lists, got (.+?) and (.+?)","errorType":"validation","errorClass":"TypeError","httpStatus":null,"severity":"error","filePath":"jax/_src/numpy/tensor_contractions.py","lineNumber":582,"sourceCode":"    if axes > min(a_ndim, b_ndim):\n      msg = \"Number of tensordot axes (axes {}) exceeds input ranks ({} and {})\"\n      raise TypeError(msg.format(axes, a.shape, b.shape))\n    contracting_dims = tuple(range(a_ndim - axes, a_ndim)), tuple(range(axes))\n  elif isinstance(axes, (tuple, list)) and len(axes) == 2:\n    ax1, ax2 = axes\n    if isinstance(ax1, int) and isinstance(ax2, int):\n      contracting_dims = ((canonicalize_axis(ax1, a_ndim),),\n                          (canonicalize_axis(ax2, b_ndim),))\n    elif isinstance(ax1, (tuple, list)) and isinstance(ax2, (tuple, list)):\n      if len(ax1) != len(ax2):\n        msg = \"tensordot requires axes lists to have equal length, got {} and {}.\"\n        raise TypeError(msg.format(ax1, ax2))\n      contracting_dims = (tuple(canonicalize_axis(i, a_ndim) for i in ax1),\n                          tuple(canonicalize_axis(i, b_ndim) for i in ax2))\n    else:\n      msg = (\"tensordot requires both axes lists to be either ints, tuples or \"\n             \"lists, got {} and {}\")\n      raise TypeError(msg.format(ax1, ax2))\n  else:\n    msg = (\"tensordot axes argument must be an int, a pair of ints, or a pair \"\n           \"of lists/tuples of ints.\")\n    raise TypeError(msg)\n  result = lax.dot_general(\n      a, b, (contracting_dims, ((), ())), precision=precision,\n      preferred_element_type=preferred_element_type,\n      out_sharding=out_sharding)\n  return lax._convert_element_type(result, preferred_element_type, output_weak_type)\n\n\n\n@export\n@api.jit(static_argnames=('precision', 'preferred_element_type'), inline=True)\ndef inner(\n    a: ArrayLike, b: ArrayLike, *, precision: lax.PrecisionLike = None,\n    preferred_element_type: DTypeLike | None = None,\n) -> Array:","sourceCodeStart":564,"sourceCodeEnd":600,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/numpy/tensor_contractions.py#L564-L600","documentation":"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).","triggerScenarios":"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.","commonSituations":"Dynamically constructed axes arguments where one branch yields an int and the other a list; passing numpy arrays or strings as axes.","solutions":["Normalize both elements to the same type (both ints or both lists/tuples of ints)","Convert numpy ints or 0-d arrays to Python int before passing"],"exampleFix":"// before\njnp.tensordot(a, b, axes=(0, [1]))\n// after\njnp.tensordot(a, b, axes=([0],[1]))","handlingStrategy":"validation","validationCode":"ax1, ax2 = axes\nassert (isinstance(ax1, int) and isinstance(ax2, int)) or (isinstance(ax1,(list,tuple)) and isinstance(ax2,(list,tuple)))","typeGuard":"def valid_axes_pair(ax1, ax2):\n    return (type(ax1) is int and type(ax2) is int) or (isinstance(ax1,(list,tuple)) and isinstance(ax2,(list,tuple)))","tryCatchPattern":null,"preventionTips":["Coerce numpy scalars to int with int() before passing axes"],"tags":["jax","tensordot","type-error","axes"],"backgroundTag":"invalid-axes-argument","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}