jax_matmul¶
- probly.representation.jax_functions.jax_matmul(x1: JaxArrayLike, x2: JaxArrayLike, /) jax.Array[source]¶
- probly.representation.jax_functions.jax_matmul(x1: JaxOperand, x2: JaxOperand, /) object
Multiply matrices or vectors, mirroring
jax.numpy.matmul.- Parameters:
x1 – Left matrix or vector operand.
x2 – Right matrix or vector operand.
- Returns:
The matrix product with broadcast batch dimensions for native operands, or the result supplied by a custom override.