Skip to content
Easy Framework

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