JAX BACKEND: add regularizer

[bonneted]

Add regularizer for JAX BACKEND.

For Jax adding the regularization is straightforward: you apply the regularization function to the nn_params and add it to the loss. see google/flax#1654

I tried to stick to TensorFlow and paddle architecture: regularizer is a net property init from the regularization parameter. (even though regularizer is not used in the net here)
In that case, the regularizer is simply the function to be applied to the net parameters to compute the loss in model.py

[lululxvi]

Why the build failed?

[bonneted]

According to github copilot:
"""
The job failed due to an error while building the dm-tree package. The specific issue is that the download of pybind11 from the specified URL failed with the status code 52, indicating "Server returned nothing (no headers, no data)".

Suggested Solution
To resolve this issue, you can try the following steps:

Check the URL: Ensure that the URL for pybind11 is correct and accessible. If the URL is outdated, update it to the latest version.

Update dm-tree: Ensure that you are using the latest version of the dm-tree package, as it may have fixed the issue with the dependency.

Retry the Build: Sometimes, transient network issues can cause such errors. Retrying the build might resolve the issue.

Modify the Workflow File: If the issue persists, modify the workflow to include the installation of pybind11 before building dm-tree.
"""

So this is probably a network issue during the build

[bonneted]

I've made and reverted a dummy commit to trigger the build again: it worked this time