jinns 0.8.1__tar.gz → 0.8.2__tar.gz

{jinns-0.8.1 → jinns-0.8.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: jinns
-Version: 0.8.1
+Version: 0.8.2
 Summary: Physics Informed Neural Network with JAX
 Author-email: Hugo Gangloff <hugo.gangloff@inrae.fr>, Nicolas Jouvin <nicolas.jouvin@inrae.fr>
 Maintainer-email: Hugo Gangloff <hugo.gangloff@inrae.fr>, Nicolas Jouvin <nicolas.jouvin@inrae.fr>

{jinns-0.8.1 → jinns-0.8.2}/doc/source/index.rst

@@ -8,6 +8,10 @@ Welcome to jinn's documentation!
 
 Changelog:
 
+* v0.8.2:
+
+    - Fix a bug: it was not possible to jit a reloaded HyperPINN model
+
 * v0.8.1:
 
     - New feature: `save_pinn` and `load_pinn` in `jinns.utils` for pre-trained

{jinns-0.8.1 → jinns-0.8.2}/jinns/utils/_hyperpinn.py

@@ -39,8 +39,8 @@ class HYPERPINN(PINN):
     static_hyper: eqx.Module
     hyperparams: list = eqx.field(static=True)
     hypernet_input_size: int
-    pinn_params_sum: ArrayLike
-    pinn_params_cumsum: ArrayLike
+    pinn_params_sum: ArrayLike = eqx.field(static=True)
+    pinn_params_cumsum: ArrayLike = eqx.field(static=True)
 
     def __init__(
         self,

{jinns-0.8.1 → jinns-0.8.2}/jinns.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: jinns
-Version: 0.8.1
+Version: 0.8.2
 Summary: Physics Informed Neural Network with JAX
 Author-email: Hugo Gangloff <hugo.gangloff@inrae.fr>, Nicolas Jouvin <nicolas.jouvin@inrae.fr>
 Maintainer-email: Hugo Gangloff <hugo.gangloff@inrae.fr>, Nicolas Jouvin <nicolas.jouvin@inrae.fr>

{jinns-0.8.1 → jinns-0.8.2}/tests/save_load_tests/test_saving_loading_hyperpinn.py

@@ -9,7 +9,7 @@ from jinns.utils import save_pinn, load_pinn
 
 
 @pytest.fixture
-def save_reload():
+def save_reload(tmpdir):
     jax.config.update("jax_enable_x64", False)
     key = random.PRNGKey(2)
 
@@ -73,7 +73,7 @@ def save_reload():
     }
 
     # Save
-    filename = "./test"
+    filename = str(tmpdir.join("test"))
     kwargs_creation = {
         "key": subkey,
         "eqx_list": eqx_list,
@@ -107,3 +107,24 @@ def test_equality_save_reload(save_reload):
         v_u_reloaded(test_points[:, 0:1], test_points[:, 1:3], params_reloaded),
         atol=1e-3,
     )
+
+
+def test_jitting_reloaded_hyperpinn(save_reload):
+    """
+    This test ensures that the reloaded hyperpinn is jit-able.
+    Some conversion of onp.array nodes can arise when reloading.
+    See this MR : https://gitlab.com/mia_jinns/jinns/-/merge_requests/32
+    jinns v0.8.2 uses eqx.field(static=True) to solve the problem.
+    This tests is here for testimony.
+    """
+
+    key, _, _, params_reloaded, u_reloaded = save_reload
+
+    v_u_reloaded = jax.vmap(
+        u_reloaded, (0, 0, {"nn_params": None, "eq_params": {"D": 0, "r": 0}})
+    )
+    v_u_reloaded_jitted = jax.jit(v_u_reloaded)
+
+    key, subkey = jax.random.split(key, 2)
+    test_points = jax.random.normal(subkey, shape=(10, 5))
+    v_u_reloaded_jitted(test_points[:, 0:1], test_points[:, 1:3], params_reloaded)

{jinns-0.8.1 → jinns-0.8.2}/tests/save_load_tests/test_saving_loading_pinn.py

@@ -9,7 +9,7 @@ from jinns.utils import save_pinn, load_pinn
 
 
 @pytest.fixture
-def save_reload():
+def save_reload(tmpdir):
     jax.config.update("jax_enable_x64", False)
     key = random.PRNGKey(2)
     eqx_list = [
@@ -28,7 +28,7 @@ def save_reload():
     params = {"nn_params": params, "eq_params": {}}
 
     # Save
-    filename = "./test"
+    filename = str(tmpdir.join("test"))
     kwargs_creation = {
         "key": subkey,
         "eqx_list": eqx_list,
@@ -57,3 +57,22 @@ def test_equality_save_reload(save_reload):
         v_u_reloaded(test_points[:, 0:1], test_points[:, 1:], params_reloaded),
         atol=1e-3,
     )
+
+
+def test_jitting_reloaded_pinn(save_reload):
+    """
+    This test ensures that the reloaded pinn is jit-able.
+    Some conversion of onp.array nodes can arise when reloading.
+    See this MR : https://gitlab.com/mia_jinns/jinns/-/merge_requests/32
+    jinns v0.8.2 uses eqx.field(static=True) to solve the problem.
+    This tests is here for testimony.
+    """
+
+    key, _, _, params_reloaded, u_reloaded = save_reload
+
+    key, subkey = jax.random.split(key, 2)
+    test_points = jax.random.normal(subkey, shape=(10, 2))
+    v_u_reloaded = jax.vmap(u_reloaded, (0, 0, None))
+    v_u_reloaded_jitted = jax.jit(v_u_reloaded)
+
+    v_u_reloaded_jitted(test_points[:, 0:1], test_points[:, 1:3], params_reloaded)

{jinns-0.8.1 → jinns-0.8.2}/tests/save_load_tests/test_saving_loading_spinn.py

@@ -9,7 +9,7 @@ from jinns.utils import save_pinn, load_pinn
 
 
 @pytest.fixture
-def save_reload():
+def save_reload(tmpdir):
     jax.config.update("jax_enable_x64", False)
     key = random.PRNGKey(2)
     d = 2
@@ -30,7 +30,7 @@ def save_reload():
     params = {"nn_params": params, "eq_params": {}}
 
     # Save
-    filename = "./test"
+    filename = str(tmpdir.join("test"))
     kwargs_creation = {
         "key": subkey,
         "d": d,
@@ -58,3 +58,21 @@ def test_equality_save_reload(save_reload):
         u_reloaded(test_points[:, 0:1], test_points[:, 1:], params_reloaded),
         atol=1e-3,
     )
+
+
+def test_jitting_reloaded_spinn(save_reload):
+    """
+    This test ensures that the reloaded spinn is jit-able.
+    Some conversion of onp.array nodes can arise when reloading.
+    See this MR : https://gitlab.com/mia_jinns/jinns/-/merge_requests/32
+    jinns v0.8.2 uses eqx.field(static=True) to solve the problem.
+    This tests is here for testimony.
+    """
+
+    key, _, _, params_reloaded, u_reloaded = save_reload
+
+    key, subkey = jax.random.split(key, 2)
+    test_points = jax.random.normal(subkey, shape=(10, 2))
+
+    u_reloaded_jitted = jax.jit(u_reloaded.__call__)
+    u_reloaded_jitted(test_points[:, 0:1], test_points[:, 1:], params_reloaded)