All Katas - ML Katas

Batch-wise Matrix Transposition

this year easy (<10 mins) | linear algebra einops tensor manipulation

### Description Given a batch of matrices, transpose each matrix in the batch. The input tensor has a shape of `(B, H, W)`, and the output should be `(B, W, H)`. ### Starter Code ```python import...

Grouped Matrix Multiplication

this year medium (<10 mins) | linear algebra tensor manipulation einsum

### Description Perform matrix multiplication on groups of matrices within a batch. For an input tensor of shape `(B, G, M, K)` and another of shape `(B, G, K, N)`, the output should be of shape...

Extract Diagonal from a Batch of Matrices

this year hard (<10 mins) | linear algebra einsum

### Description Given a batch of square matrices of shape `(B, N, N)`, extract the diagonal of each matrix. The output should be a tensor of shape `(B, N)`. This can be achieved with...

Matrix Multiplication Basics

this year easy (<10 mins) | linear algebra tensors torch.mm

Implement a function in PyTorch that multiplies two matrices using `torch.mm`. ### Problem Write a function `matmul(A, B)` that takes two 2D tensors `A` and `B` and returns their matrix product. -...