jax.lax.batch_matmul

jax.lax.batch_matmul(lhs, rhs, precision=None)

Batch matrix multiplication.

Parameters:

• lhs (Array)

• rhs (Array)

• precision (Union[str, Precision, tuple[str, str], tuple[Precision, Precision], None])

Return type:

Array