def contract(
    names: str | Sequence[named_axes.AxisName],
    left: NamedArray,
    right: NamedArray,
) -> NamedArray:
  """Contracts two named arrays along the given axis names.

  Args:
    names: The axis names to contract. Can also be a single string axis name.
    left: The left-hand side of the contraction, as a NamedArray with only named
      axes.
    right: The right-hand side of the contraction, as a NamedArray with only
      named axes.

  Returns:
    Result of the contraction, which will have all names present in either
    `left` or `right`, except for the names in `names` which will be contracted
    away.
  """
  names = (names,) if isinstance(names, str) else names
  return named_axes.nmap(jnp.tensordot)(
      left.untag(*names), right.untag(*names), axes=len(names)
  )
Contracts two named arrays along the given axis names.

  Args:
    names: The axis names to contract. Can also be a single string axis name.
    left: The left-hand side of the contraction, as a NamedArray with only named
      axes.
    right: The right-hand side of the contraction, as a NamedArray with only
      named axes.

  Returns:
    Result of the contraction, which will have all names present in either
    `left` or `right`, except for the names in `names` which will be contracted
    away.
  
penzai/nn/linear_and_affine.py