Merge pull request #75 from markusschmitt/jax_flax_version_compatibility

Compatibility with newest Jax and Flax versions

jVMC/__init__.py

@@ -10,5 +10,5 @@
 from .version import __version__
 from .global_defs import set_pmap_devices
 
-from jax.config import config
-config.update("jax_enable_x64", True)
+import jax
+jax.config.update("jax_enable_x64", True)

jVMC/nets/cnn.py

@@ -1,6 +1,5 @@
 import jax
-from jax.config import config
-config.update("jax_enable_x64", True)
+jax.config.update("jax_enable_x64", True)
 import flax
 import flax.linen as nn
 import jax.numpy as jnp

jVMC/nets/initializers.py

@@ -1,6 +1,5 @@
 import jax
-from jax.config import config
-config.update("jax_enable_x64", True)
+jax.config.update("jax_enable_x64", True)
 import flax
 import flax.linen as nn
 

jVMC/nets/rbm.py

@@ -1,6 +1,5 @@
 import jax
-from jax.config import config
-config.update("jax_enable_x64", True)
+jax.config.update("jax_enable_x64", True)
 import flax
 #from flax import nn
 import flax.linen as nn

jVMC/nets/rnn1d_general.py

@@ -1,6 +1,5 @@
 import jax
-from jax.config import config
-config.update("jax_enable_x64", True)
+jax.config.update("jax_enable_x64", True)
 import flax
 import flax.linen as nn
 import numpy as np
@@ -119,10 +118,10 @@ def setup(self):
             if self.cell == "RNN":
                 self.cells = [RNNCell(actFun=self.actFun, initFun=self.initFunction, dtype=self.dtype) for _ in range(self.depth)]
             elif self.cell == "LSTM":
-                self.cells = [LSTMCell() for _ in range(self.depth)]
+                self.cells = [LSTMCell(features=self.hiddenSize) for _ in range(self.depth)]
                 self.zero_carry = jnp.zeros((self.depth, 2, self.hiddenSize), dtype=self.dtype)
             elif self.cell == "GRU":
-                self.cells = [GRUCell() for _ in range(self.depth)]
+                self.cells = [GRUCell(features=self.hiddenSize) for _ in range(self.depth)]
             else:
                 ValueError("Cell name not recognized.")
         else:
@@ -179,16 +178,20 @@ def rnn_cell_sample(self, carry, x):
 
 
 class GRUCell(nn.Module):
+    features: int
+
     @nn.compact
     def __call__(self, carry, state):
-        current_carry, newR = nn.GRUCell(**init_fn_args(recurrent_kernel_init=jax.nn.initializers.orthogonal(dtype=global_defs.tReal)))(carry, state)
+        current_carry, newR = nn.GRUCell(features=self.features, **init_fn_args(recurrent_kernel_init=jax.nn.initializers.orthogonal(dtype=global_defs.tReal)))(carry, state)
         return current_carry, newR[0]
 
 
 class LSTMCell(nn.Module):
+    features: int
+
     @nn.compact
     def __call__(self, carry, state):
-        current_carry, newR = nn.OptimizedLSTMCell(**init_fn_args(recurrent_kernel_init=jax.nn.initializers.orthogonal(dtype=global_defs.tReal)))(carry, state)
+        current_carry, newR = nn.OptimizedLSTMCell(features=self.features, **init_fn_args(recurrent_kernel_init=jax.nn.initializers.orthogonal(dtype=global_defs.tReal)))(carry, state)
         return jnp.asarray(current_carry), newR
 
 

jVMC/nets/rnn2d_general.py

@@ -1,6 +1,5 @@
 import jax
-from jax.config import config
-config.update("jax_enable_x64", True)
+jax.config.update("jax_enable_x64", True)
 # config.update('jax_disable_jit', True)
 import flax
 import flax.linen as nn
@@ -127,7 +126,7 @@ def setup(self):
             if self.cell == "RNN":
                 self.cells = [RNNCell(actFun=self.actFun, initFun=self.initFunction, dtype=self.dtype) for _ in range(self.depth)]
             elif self.cell == "LSTM":
-                self.cells = [LSTMCell() for _ in range(self.depth)]
+                self.cells = [LSTMCell(features=self.hiddenSize) for _ in range(self.depth)]
                 self.zero_carry = jnp.zeros((self.L, self.depth, 2, self.hiddenSize), dtype=self.dtype)
             elif self.cell == "GRU":
                 self.cells = [GRUCell() for _ in range(self.depth)]
@@ -241,6 +240,8 @@ def rnn_cell_H_sample(self, carry, x):
 
 
 class GRUCell(nn.Module):
+    features: int
+
     @nn.compact
     def __call__(self, carryH, carryV, state):
         cellCarryH = nn.Dense(features=carryH.shape[-1],
@@ -249,12 +250,14 @@ def __call__(self, carryH, carryV, state):
         cellCarryV = nn.Dense(features=carryV.shape[-1],
                               use_bias=False,
                               dtype=global_defs.tReal)
-        current_carry, newR = nn.GRUCell(param_dtype=global_defs.tReal)(cellCarryH(carryH) + cellCarryV(carryV), state)
+        current_carry, newR = nn.GRUCell(features=self.features, param_dtype=global_defs.tReal)(cellCarryH(carryH) + cellCarryV(carryV), state)
 
         return current_carry, newR[0]
 
 
 class LSTMCell(nn.Module):
+    features: int
+
     @nn.compact
     def __call__(self, carryH, carryV, state):
         cellCarryH = nn.Dense(features=carryH.shape[-1],
@@ -263,7 +266,7 @@ def __call__(self, carryH, carryV, state):
         cellCarryV = nn.Dense(features=carryV.shape[-1],
                               use_bias=False,
                               param_dtype=global_defs.tReal)
-        current_carry, newR = nn.OptimizedLSTMCell(param_dtype=global_defs.tReal)(cellCarryH(carryH) + cellCarryV(carryV), state)
+        current_carry, newR = nn.OptimizedLSTMCell(features=self.features, param_dtype=global_defs.tReal)(cellCarryH(carryH) + cellCarryV(carryV), state)
 
         return jnp.asarray(current_carry), newR
 

jVMC/util/util.py

@@ -1,6 +1,5 @@
 import jax
-from jax.config import config
-config.update("jax_enable_x64", True)
+jax.config.update("jax_enable_x64", True)
 
 import jax.numpy as jnp
 import jax.random as random

jVMC/version.py

@@ -1,2 +1,2 @@
 """Current jVMC version at head on Github."""
-__version__ = "1.3.1"
+__version__ = "1.4.0"

jVMC/vqs.py

@@ -1,6 +1,5 @@
 import jax
-from jax.config import config
-config.update("jax_enable_x64", True)
+jax.config.update("jax_enable_x64", True)
 from jax import jit, grad, vmap
 from jax import numpy as jnp
 from jax import random
@@ -481,8 +480,13 @@ def params(self):
     @params.setter
     def params(self, val):
         # Replace 'params' in parameters by `val`
-        self.parameters = freeze({
-                                **unfreeze(self.parameters.pop("params")[0]),
-                                "params": unfreeze(val)
-                                })
+        ps = unfreeze(self.parameters)
+        ps["params"] = unfreeze(val)
+        if isinstance(self.parameters, flax.core.frozen_dict.FrozenDict):
+            ps = freeze(ps)
+        self.parameters = ps
+        # self.parameters = freeze({
+        #                        **unfreeze(self.parameters.pop("params")[0]),
+        #                        "params": unfreeze(val)
+        #                        })
 

setup.py

@@ -7,7 +7,7 @@
 with open("README.md", "r") as fh:
     long_description = fh.read()
 
-DEFAULT_DEPENDENCIES = ["setuptools", "wheel", "numpy", "openfermion", "jax>=0.4.1,<=0.4.20", "jaxlib>=0.4.1,<=0.4.20", "flax>=0.6.4,<=0.6.11", "mpi4py", "h5py", "PyYAML", "matplotlib", "scipy<1.13"] # Scipy version restricted, because jax is currently incompatible with new function namespace scipy.sparse.tril
+DEFAULT_DEPENDENCIES = ["setuptools", "wheel", "numpy", "openfermion", "jax>=0.4.12,<=0.4.31", "jaxlib>=0.4.12,<=0.4.31", "flax>=0.7.0", "mpi4py", "h5py", "PyYAML", "matplotlib", "scipy<1.13"] # Scipy version restricted, because jax is currently incompatible with new function namespace scipy.sparse.tril
 #CUDA_DEPENDENCIES = ["setuptools", "wheel", "numpy", "jax[cuda]>=0.2.11,<=0.2.25", "flax>=0.3.6,<=0.3.6", "mpi4py", "h5py"]
 DEV_DEPENDENCIES = DEFAULT_DEPENDENCIES + ["sphinx", "mock", "sphinx_rtd_theme", "pytest", "pytest-mpi"]
 

tests/mpi_wrapper_test.py

@@ -1,8 +1,7 @@
 import unittest
 
 import jax
-from jax.config import config
-config.update("jax_enable_x64", True)
+jax.config.update("jax_enable_x64", True)
 
 import jax.numpy as jnp
 

tests/nets_test.py

@@ -1,12 +1,9 @@
 import unittest
-import sys
-sys.path.append(sys.path[0] + "../..")
 import jVMC
 import jVMC.nets as nets
 
 import jax
-from jax.config import config
-config.update("jax_enable_x64", True)
+jax.config.update("jax_enable_x64", True)
 import jax.random as random
 import jax.numpy as jnp
 import numpy as np

tests/operator_test.py

@@ -1,8 +1,7 @@
 import unittest
 
 import jax
-from jax.config import config
-config.update("jax_enable_x64", True)
+jax.config.update("jax_enable_x64", True)
 
 import jax.random as random
 import jax.numpy as jnp

tests/povm_test.py

@@ -51,8 +51,8 @@ def copy_dict(a):
         self.stepper = jVMC.util.stepper.Heun(timeStep=dt)  # ODE integrator
 
     def test_matrix_to_povm(self):
-        unity = jnp.eye(2)
-        zero_matrix = jnp.zeros((2, 2))
+        unity = jnp.eye(2, dtype=jVMC.global_defs.tCpx)
+        zero_matrix = jnp.zeros((2, 2), dtype=jVMC.global_defs.tCpx)
 
         system_data = {"dim": "1D", "L": 2}
         povm = jVMC.operator.POVM(system_data)
@@ -78,8 +78,8 @@ def test_matrix_to_povm(self):
         self.assertRaises(ValueError, op.matrix_to_povm, zero_matrix, povm.M, povm.T_inv, mode='wrong_mode')
 
     def test_adding_operator(self):
-        unity = jnp.eye(2)
-        zero_matrix = jnp.zeros((2, 2))
+        unity = jnp.eye(2, dtype=jVMC.global_defs.tCpx)
+        zero_matrix = jnp.zeros((2, 2), dtype=jVMC.global_defs.tCpx)
 
         system_data = {"dim": "1D", "L": 2}
         povm = jVMC.operator.POVM(system_data)

tests/sampler_test.py

@@ -1,8 +1,7 @@
 import unittest
 
 import jax
-from jax.config import config
-config.update("jax_enable_x64", True)
+jax.config.update("jax_enable_x64", True)
 import jax.random as random
 import jax.numpy as jnp
 

tests/symmetries_test.py

@@ -1,8 +1,7 @@
 import unittest
 
 import jax
-from jax.config import config
-config.update("jax_enable_x64", True)
+jax.config.update("jax_enable_x64", True)
 import jax.numpy as jnp
 
 import numpy as np

tests/tdvp_test.py

@@ -1,8 +1,7 @@
 import unittest
 
 import jax
-from jax.config import config
-config.update("jax_enable_x64", True)
+jax.config.update("jax_enable_x64", True)
 
 import jax.random as random
 import jax.numpy as jnp