`from_torch` and `to_torch` don't work for functions on mixed function spaces

[jshipton]

On Slack @connorjward suggested that the fix might be straightforward... I'd be happy to help but would need some guidance...

[connorjward]

I think that the simplest way to achieve this is to replace:

x_F = Function(V, val=x.detach().numpy())

with

x_F = Function(V)
with x_F.dat.vec_wo as vec:
    vec.buffer[...] = x.detach().numpy()

here.

For to_torch you probably want

x_P = torch.tensor(np.ravel(x.dat.data_ro), **kwargs)

to turn into

with x.dat.vec_ro as vec:
    x_P = torch.tensor(np.ravel(vec.buffer_r), **kwargs)

This is all necessary because a mixed Function stores its data as separate numpy arrays per field. Using the vec context managers concatenates the arrays.

If you could add a test and maybe also do the same thing for JAX that would be really cool.