Flax/nnx backend

[frazane]

Description

Introduces flax / nnx as backend for the library.

Tasks

• [x] pass all standard tests
• [x] pass integration tests
• [x] check and eventually update all docstrings
• [x] pass docstrings tests
• [x] run all documentation notebooks
• [x] adapt markdown cells in documentation notebooks
• [x] new docs pages

Running documentation notebooks

• [x] barycentres
• [x] bayesian_optimisation
• [x] classification
• [x] collapsed_vi
• [x] constructing_new_kernels
• [x] decision_making
• [x] deep_kernels
• [x] graph_kernels
• [x] intro_to_gps
• [x] intro_to_kernels
• [x] likelihoods_guide
• [x] oceanmodelling
• [x] poisson
• [x] regression
• [x] uncollapsed_vi
• [x] utils
• [x] yacht