def _non_batched_matmul(lhs, rhs, lhs_contraction, rhs_contraction):
  """Compute the non-batched matrix multiplication.

  If it is the general non-batched/single-batched matrix multiplication,
  use the highly optimized kernel `tf.tensordot` to handle it.

  Args:
    lhs: an array (the left-hand side matrix/vector to be multiplied)
    rhs: an array (the right-hand side matrix/vector to be multiplied)
    lhs_contraction: Sequence[int] (the contraction dimensions of lhs)
    rhs_contraction: Sequence[int] (the contraction dimensions of rhs)

  Returns:
    An array that contains the result.
  """
  return tf.tensordot(
      lhs, rhs, axes=(list(lhs_contraction), list(rhs_contraction)))
Compute the non-batched matrix multiplication.

  If it is the general non-batched/single-batched matrix multiplication,
  use the highly optimized kernel `tf.tensordot` to handle it.

  Args:
    lhs: an array (the left-hand side matrix/vector to be multiplied)
    rhs: an array (the right-hand side matrix/vector to be multiplied)
    lhs_contraction: Sequence[int] (the contraction dimensions of lhs)
    rhs_contraction: Sequence[int] (the contraction dimensions of rhs)

  Returns:
    An array that contains the result.
  
trax/tf_numpy/extensions/extensions.py