Sparse Linear Layer

Hi, I want to create a sparse linear layer, where I mask certain weights & gradients according to a certain mask. Is there an example where this is done efficiently?

Do you mean something like: