ML Katas

Concatenate Tensors Along a New Axis

easy (<10 mins) einops tensor manipulation
this year by E

Description

Given a list of tensors of the same shape, concatenate them along a new axis. For example, given 3 tensors of shape (H, W), you want to create a single tensor of shape (3, H, W).

Starter Code

import torch
from einops import stack

def stack_tensors(tensor_list):
    # Your einops code here
    pass

Verification

Create a list of 5 tensors, each of shape (10, 20). The output should be a single tensor of shape (5, 10, 20).