Upgrade JAX, TensorFlow + Keras, TFP, and typing_extensions for Python 3.13 compatibility

[GabrielOduori]

PR Summary (one-paragraph)

This pull request updates JAX imports, standardizes type annotations, and upgrades dependencies to ensure full compatibility with Python 3.13. It replaces legacy DeviceArray annotations with Array from JAX, enables 64-bit computation via jax.config.update("jax_enable_x64", True), and upgrades JAX, TensorFlow (with TF-Keras), TensorFlow Probability, GPflow, and typing_extensions to versions compatible with Python 3.13. These changes allow BayesNewton tests to run correctly under Python 3.13, while dataset-dependent tests may still fail if the corresponding experimental data is missing.

---

Detailed changes

1. JAX import update

   • Changed from:

     # from jax.config import config
     # config.update("jax_enable_x64", True)

     to:

     import jax
     jax.config.update("jax_enable_x64", True)

   • Ensures proper configuration for 64-bit computations.

2. Type annotations updated

   • Changed:

     # Y: np.DeviceArray
     # ind: np.DeviceArray

     to:

     Y: Array
     ind: Array

   • Aligns type annotations with JAX’s current array types and avoids deprecation warnings.

3. Dependencies upgraded

   • jax and jaxlib: 0.4.14 → 0.4.34

   • tensorflow: upgraded to 2.20.0

   • tf-keras: upgraded to 2.20.1

   • tensorflow-probability: upgraded to 0.25.0

   • gpflow: upgraded to 2.9.2

   • typing_extensions: upgraded to 4.15.0

   • Rationale: fixes Python 3.13 compatibility issues, including:

     • TypeError: type 'typing.TypeVar' is not an acceptable base type
     • TF-Keras import checks in TensorFlow Probability
     • JAX 64-bit computation support

4. Impact on tests

   • Tests now run without TypeVar errors.
   • Dataset-dependent tests (e.g., test_vs_gpflow_shutters.py) may fail if experiment data is missing, unrelated to these changes.