Torch backend

[BenjaminMidtvedt]

This is the proposal for enabling the use of a torch backend. The class / method names are not at all set in stone. In fact, I would very much appreciate feedback.

The core idea is as follows:

Internally, we use a special xp module instead of numpy. xp is based on array-api-compat, which is python library that provided a unified api to operate on arrays, based on the array-api standard (which is what numpy uses, and is the most established standard).

xp provides two core functionalities for us. First, as mentioned, it provides a unified interface for torch and numpy arrays (as well as cupy, dask and jax). This means that as long as we use xp everywhere, we will be cross-compatible. We do not have to change the code path based on the array provider.

Second, I've set it up to automatically create arrays of the type defined by a backend options in the deeptrack config object. So, if you call config.set_backend_numpy(), then xp.zeros((100, 100)) will create a numpy array, and if you call config.set_backend_torch(), the same code will create a torch tensor.

Feature.numpy(), Feature.torch()

there are two methods, f.numpy() and f.torch(), these can be used to override the global backend choice for a specific pipieline. This means that you can use different backends in different parts of the full pipeline. Note, however, that there is no automatic conversion between the backend (passing a tensor to a numpy feature will give an error). This is a choice made for simplicity and performance.

Device

xp array creation can be provided a device argument. For numpy, the only options are "cpu" and None, while for torch, all device options are available (both strings, or torch.device objects). I propose that each feature has a device flag that is set similarly to the backend. We will have to ensure that any array creation calls pass this flag: xp.zeros((100, 100), device=self.device)).

Again, I propose that we do not do automatic device transfers, and instead leave this to the user. However, we could provide a simple method to do so on the Feature class:

foo = foo_pipeline.torch().cuda()
bar = bar_pipeline.numpy().cpu()
foobar = foo.output(device="cpu", backend="numpy") >> bar

Examples

I have in this branch two examples in the root. These should be removed or moved to a better place before the PR is accepted. They demonstrate a minimal example of how this works, and how the torch backend could be used in a gradient descent setting.

Testing

A good starting point for testing this new backend would be to run the existing suite of tests with the torch backend.
I propose that each feature is additionally tested to ensure the output array type is of the right backend.
I also propose that each feature is additionally tested to ensure that gradients flow through the object. This should include all numeric properties as well. Full gradient compatibility is likely to take more time, and I propose that full compatibility is not required on first release.

Users

Users are not required to use xp for their own features. It is only required for cross-compatibility between backends. As such, old features written using numpy will still work identically.