Torch backend implementation

[MUCDK]

Description of feature

Once we merged #228, #235, #239, #240, we can start with the torch backend implementation. Therefore

• we will keep one CellFlow class for both backends
• have a src/backends/jax and a src/backends/torch directory
• CellFlow.prepare_model will have an argument backend: Literal["jax", "torch"]
• for now, let's not implement GENOT for torch, let's just go with OTFlowMatching, to keep things simple-
• the only problem which arises is that prepare_model takes backend-specific arguments, namely match_fn, optimizer, and vf_act_fn.

For the last point, I see the following as the best solution
allow passing both jax and torch instances, setting per default None , describe the default in the docs, and eventually instantiating them later in the solver classes

[MUCDK]

Pinging @AlejandroTL to track progress

[MUCDK]

Two options:

1. directories along the lines of cellflow/src/backend and cellflow/src/base (backend-dependent modules higher up)
2. directories along the lines of cellflow/src/networks/backends/jax cellflow/src/models/backends/jax (backend-dependent modules further down)

@selmanozleyen any opinion?