Matrix Multiplication

Multiply two matrices using the framework’s built-in matrix multiplication.

Input:

• A: A 2D tensor of shape (m, k)
• B: A 2D tensor of shape (k, n)

Output: A 2D tensor of shape (m, n) equal to A @ B.

API Reference:

• PyTorch: torch.matmul(A, B) or A @ B
• JAX: jnp.matmul(A, B) or A @ B