jax.lax.batch_matmul

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

Batch matrix multiplication.

Parameters:

• lhs (Array)

• rhs (Array)

• precision (PrecisionLike | None)

Return type:

Array