Modify custom `torch.autograd.Function`s to use `setup_context` instead of passing `ctx` as an argument to `forward`

[tvercaut]

See https://pytorch.org/docs/master/notes/extending.func.html This should help make our code compatible with vmap

[theo-barfoot]

I have given this a go. See: SparseMatMul on the batched-mm-vmap branch.

This implementation only works for the forward pass of the COO batched mm.

It fails for:

• COO backward: [undefined]NotImplementedError: Could not run 'aten::select_backward' with arguments from the 'SparseCPU' backend
• CSR foward: Failed: [undefined]RuntimeError: Sparse CSR tensors do not have strides
• CSR backward: [undefined]RuntimeError: Sparse CSR tensors do not have strides

The CSR methods fails (at least in the first instance) due to the function _add_batch_dim not supporting CSR tensors.