Add support for Jax / Haiku

Goal : implement HaikuModel and JaxNumpyVectors.

Must have

• Do MVP
  • Create JaxNumpyVector
  • Create HaikyModel, with a tax version of all abstract methods
  • Create a minimal example for testing
  • Deal with flattening and unflattening params, using jax.tree_utils
• Refine MVP
  • Add clipping : use optax-inspired clipping per sample
  • Add jit
  • Add GPU support
  • Add proper randomization
  • Add typing
  • Use the new DataInfoFIeld
  • Properly document
  • Rebase and lint
  • Full test suite
  • Add dependencies
• Examine edge cases
  • Deal with statefulness in haiku modules (Implement LSTM with unroll_net, add RNN to test)
  • Take a closer look at multi-input models (data_type at init, inputs as list ) : use unpack_batch
  • Look at frozen weights (don't forget to update test, see 'MLP-tune')
  • Figure out why per-sample gradient calculations are so slow.
  • Ensure BatchNorm works

Nice to haves :

• Upgrade my JAX :
  • Reintroduce the State class to shorten method signatures, using @chex.dataclass
  • Make custom class to @jit whole methods while using instance attributes (source):
  • Allow for external PRNGSequence initilaization, to allow optimization via the reserve argument (source)
  • explore using hk.Initializer
  • Add @chex.variant to testing
• Enforce checks on model and loss function passed as input to HaikuModel
• Refine loss :
  • Accept optax losses, do the mapping on loss inputs (inspect.get_signature)
  • Add kwargs to _compute_loss and loss_function to cover all cases of optax loss
  • Look at mixed precision
• Refine serialization : avoid using pickling.

declearn/model/haiku/utils/__init__.py

+# coding: utf-8
+
+# Copyright 2023 Inria (Institut National de Recherche en Informatique
+# et Automatique)
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Utils for jax backend support code.
+
+GPU/CPU backing device management utils:
+* select_device:
+    Select a backing device to use based on inputs and availability.
+"""
+
+from ._gpu import select_device