updating versions

Author: lenarttreven

bsm/bayesian_regression/bayesian_neural_networks/bnn.py

@@ -5,6 +5,7 @@
 from typing import Sequence
 
 import chex
+import flax
 import jax
 import jax.numpy as jnp
 import jax.random as jr
@@ -144,7 +145,7 @@ def step_jit(self,
     def _init(self, key):
         variables = self.model.init(key, jnp.ones(shape=(self.input_dim,)))
         if 'params' in variables:
-            stats, params = variables.pop('params')
+            stats, params = flax.core.pop(variables, 'params')
         else:
             stats, params = variables
         del variables  # Delete variables to avoid wasting resources

bsm/bayesian_regression/bayesian_neural_networks/deterministic_ensembles.py

@@ -47,7 +47,7 @@ def apply_eval(self,
 
     noise_level = 0.1
     d_l, d_u = 0, 10
-    xs = jnp.linspace(d_l, d_u, 256).reshape(-1, 1)
+    xs = jnp.linspace(d_l, d_u, 32).reshape(-1, 1)
     ys = jnp.concatenate([jnp.sin(xs), jnp.cos(xs)], axis=1)
     ys = ys + noise_level * random.normal(key=random.PRNGKey(0), shape=ys.shape)
     data_std = noise_level * jnp.ones(shape=(output_dim,))

bsm/bayesian_regression/bayesian_recurrent_neural_networks/rnn_ensembles.py

@@ -4,6 +4,7 @@
 from typing import Sequence, Dict, Tuple, Optional
 
 import chex
+import flax
 import jax
 import jax.numpy as jnp
 import matplotlib.pyplot as plt
@@ -117,7 +118,7 @@ def eval_ll(self,
     def _init(self, key):
         variables = self.model.init(key, jnp.ones(shape=(1, self.input_dim)))
         if 'params' in variables:
-            stats, params = variables.pop('params')
+            stats, params = flax.core.pop(variables, 'params')
         else:
             stats, params = variables
         del variables  # Delete variables to avoid wasting resources

bsm/bayesian_regression/gaussian_processes/__init__

@@ -17,7 +17,7 @@
 from bsm.bayesian_regression.bayesian_regression_model import BayesianRegressionModel
 from bsm.bayesian_regression.gaussian_processes.kernels import Kernel, RBF
 from bsm.utils.normal_with_aleatoric import ExtendedNormal
-from bsm.utils.normalization import Normalizer, DataStats, Data, Stats
+from bsm.utils.normalization import Normalizer, DataStats, Data
 
 
 @chex.dataclass
@@ -35,10 +35,12 @@ def __init__(self,
                  lr_rate: optax.Schedule | float = optax.constant_schedule(1e-2),
                  seed: int = 0,
                  logging_wandb: bool = True,
+                 normalize: bool = True,
                  *args,
                  **kwargs
                  ):
         super().__init__(*args, **kwargs)
+        self.normalize = normalize
         if kernel is None:
             kernel = RBF(self.input_dim)
         self.kernel = kernel
@@ -56,13 +58,16 @@ def __init__(self,
         self.kernel_multiple_output = vmap(self.kernel.apply, in_axes=(None, None, 0), out_axes=0)
 
     def init(self, key: chex.PRNGKey) -> GPModelState:
-        keys = jr.split(key, self.output_dim)
-        params = vmap(self.kernel.init)(keys)
         inputs = jnp.zeros(shape=(1, self.input_dim))
         outputs = jnp.zeros(shape=(1, self.output_dim))
         data = Data(inputs=inputs, outputs=outputs)
-        data_stats = self.normalizer.compute_stats(data.inputs)
-        return GPModelState(history=data, data_stats=data_stats, params=params)
+        if self.normalize:
+            data_stats = self.normalizer.compute_stats(data.inputs)
+        else:
+            data_stats = self.normalizer.init_stats(data.inputs)
+        keys = jr.split(key, self.output_dim)
+        params = vmap(self.kernel.init)(keys)
+        return GPModelState(params=params, data_stats=data_stats, history=data)
 
     def loss(self, vmapped_params, inputs, outputs, data_stats: DataStats):
         assert inputs.shape[0] == outputs.shape[0]
@@ -105,7 +110,10 @@ def fit_model(self,
                   model_state: GPModelState) -> GPModelState:
         vmapped_params = model_state.params
         opt_state = self.tx.init(vmapped_params)
-        data_stats = self.normalizer.compute_stats(data)
+        if self.normalize:
+            data_stats = self.normalizer.compute_stats(data)
+        else:
+            data_stats = self.normalizer.init_stats(data)
 
         def f(carry, _):
             opt_state, vmapped_params = carry

bsm/bayesian_regression/gaussian_processes/gaussian_processes.py

@@ -2,16 +2,16 @@
 import jax.numpy as jnp
 import jax.random as jr
 import optax
-from jax import vmap
 
 from bsm.bayesian_regression.bayesian_neural_networks.bnn import BNNState
 from bsm.bayesian_regression.bayesian_neural_networks.bnn import BayesianNeuralNet
 from bsm.bayesian_regression.bayesian_neural_networks.deterministic_ensembles import DeterministicEnsemble
+from bsm.bayesian_regression.bayesian_neural_networks.fsvgd_ensemble import DeterministicFSVGDEnsemble, \
+    ProbabilisticFSVGDEnsemble
 from bsm.bayesian_regression.bayesian_neural_networks.probabilistic_ensembles import ProbabilisticEnsemble
-from bsm.bayesian_regression.bayesian_neural_networks.fsvgd_ensemble import DeterministicFSVGDEnsemble, ProbabilisticFSVGDEnsemble
 from bsm.statistical_model.abstract_statistical_model import StatisticalModel
 from bsm.utils.normalization import Data
-from bsm.utils.type_aliases import StatisticalModelState, StatisticalModelOutput
+from bsm.utils.type_aliases import StatisticalModelState
 
 
 class BNNStatisticalModel(StatisticalModel[BNNState]):
@@ -91,6 +91,7 @@ def update(self, stats_model_state: StatisticalModelState[BNNState], data: Data)
         plt.plot(test_xs.reshape(-1), test_ys[:, j], label='True', color='green')
         by_label = dict(zip(labels, handles))
         plt.legend(by_label.values(), by_label.keys())
+        plt.savefig(f'bnn_{j}.pdf')
         plt.show()
 
     num_test_points = 1000
@@ -102,4 +103,4 @@ def update(self, stats_model_state: StatisticalModelState[BNNState], data: Data)
         plt.plot(in_domain_test_xs, in_domain_preds.mean[:, j], label='Mean', color='blue')
         plt.plot(in_domain_test_xs, in_domain_test_ys[:, j], label='Fun', color='Green')
         plt.legend()
-        plt.show()
+        plt.show()

bsm/statistical_model/bnn_statistical_model.py

@@ -4,10 +4,10 @@
 import optax
 from jax import vmap
 
-from bsm.statistical_model.abstract_statistical_model import StatisticalModel
 from bsm.bayesian_regression.gaussian_processes.gaussian_processes import GPModelState, GaussianProcess
+from bsm.statistical_model.abstract_statistical_model import StatisticalModel
 from bsm.utils.normalization import Data
-from bsm.utils.type_aliases import StatisticalModelState, StatisticalModelOutput
+from bsm.utils.type_aliases import StatisticalModelState
 
 
 class GPStatisticalModel(StatisticalModel[GPModelState]):
@@ -18,9 +18,11 @@ def __init__(self,
                  delta: float = 0.1,
                  num_training_steps: int = 1000,
                  beta: chex.Array | optax.Schedule | None = None,
+                 normalize: bool = True,
                  *args, **kwargs
                  ):
-        model = GaussianProcess(input_dim=input_dim, output_dim=output_dim, *args, **kwargs)
+        self.normalize = normalize
+        model = GaussianProcess(input_dim=input_dim, output_dim=output_dim, normalize=normalize, *args, **kwargs)
         super().__init__(input_dim, output_dim, model)
         self.model = model
         self.f_norm_bound = f_norm_bound
@@ -41,8 +43,11 @@ def update(self, stats_model_state: StatisticalModelState, data: Data) -> Statis
             return StatisticalModelState(model_state=new_model_state, beta=beta)
 
     def compute_beta(self, model_state: GPModelState, data: Data):
-        inputs_norm = vmap(self.model.normalizer.normalize, in_axes=(0, None))(data.inputs,
-                                                                               model_state.data_stats.inputs)
+        if self.normalize:
+            inputs_norm = vmap(self.model.normalizer.normalize, in_axes=(0, None))(data.inputs,
+                                                                                   model_state.data_stats.inputs)
+        else:
+            inputs_norm = data.inputs
         covariance_matrix = self.model.m_kernel_multiple_output(inputs_norm, inputs_norm, model_state.params)
         covariance_matrix = covariance_matrix / (self.model.output_stds ** 2)[:, None, None]
         covariance_matrix = covariance_matrix + jnp.eye(covariance_matrix.shape[-1])[None, :, :]
@@ -92,4 +97,5 @@ def compute_beta(self, model_state: GPModelState, data: Data):
         plt.plot(test_xs.reshape(-1), test_ys[:, j], label='True', color='green')
         by_label = dict(zip(labels, handles))
         plt.legend(by_label.values(), by_label.keys())
+        plt.savefig(f'gp_{j}.pdf')
         plt.show()

bsm/statistical_model/gp_statistical_model.py

@@ -30,13 +30,25 @@ def __init__(self, num_correction=1e-6):
     def compute_stats(self, data: PyTree) -> PyTree[Stats]:
         return jtu.tree_map(self.get_stats, data)
 
+    def init_stats(self, data: PyTree) -> PyTree[Stats]:
+        return jtu.tree_map(self._init_stats, data)
+
     @partial(jax.jit, static_argnums=0)
     def get_stats(self, data: chex.Array) -> Stats:
         assert data.ndim == 2
         mean = jnp.mean(data, axis=0)
-        std = jnp.std(data, axis=0) + self.num_correction
+        if data.shape[0] > 1:
+            std = jnp.std(data, axis=0) + self.num_correction
+        else:
+            std = jnp.ones_like(mean)
         return Stats(mean, std)
 
+    @partial(jax.jit, static_argnums=0)
+    def _init_stats(self, data: chex.Array) -> Stats:
+        assert data.ndim == 2
+        mean = jnp.mean(data, axis=0)
+        return Stats(jnp.zeros_like(mean), jnp.ones_like(mean))
+
     @partial(jax.jit, static_argnums=0)
     def normalize(self, datum: chex.Array, stats: Stats) -> chex.Array:
         assert datum.ndim == 1

bsm/utils/normalization.py

@@ -3,17 +3,17 @@
 from setuptools import setup, find_packages
 
 required = [
-    'flax>=0.4.1,<=0.7.0',
-    'jax==0.4.13',
+    'flax>=0.7.0',
+    'jax>=0.4.13',
     'jaxtyping>=0.2.20',
-    'jaxlib==0.4.13',
-    'pytest==7.4.0',
+    'jaxlib>==0.4.13',
+    'pytest>=7.4.0',
     'matplotlib>=3.5.1',
     'numpy>=1.22.2',
     'optax>=0.1.1',
     'scipy>=1.8.0',
     'wandb>=0.12.11',
-    'distrax~=0.1.2',
+    'distrax @ git+https://github.com/deepmind/distrax.git',
     'argparse-dataclass>=0.2.1',
     'jaxutils',
     'chex',

setup.py

@@ -1,11 +1,9 @@
 import jax.numpy as jnp
 import jax.random as jr
-from jax import vmap
 
 from bsm.bayesian_regression import DeterministicEnsemble
 from bsm.statistical_model import BNNStatisticalModel
 from bsm.utils.normalization import Data
-from bsm.utils.type_aliases import StatisticalModelOutput
 
 key = jr.PRNGKey(0)
 input_dim = 1

tests/test_deterministic_ensemble_bnn_sm.py

@@ -1,11 +1,9 @@
 import jax.numpy as jnp
 import jax.random as jr
-from jax import vmap
 
 from bsm.bayesian_regression import DeterministicFSVGDEnsemble
 from bsm.statistical_model import BNNStatisticalModel
 from bsm.utils.normalization import Data
-from bsm.utils.type_aliases import StatisticalModelOutput
 
 key = jr.PRNGKey(0)
 input_dim = 1

tests/test_deterministic_fsvg_ensemble_bnn_sm.py

@@ -3,10 +3,9 @@
 from jax import vmap
 
 from bsm.bayesian_regression import DeterministicGRUEnsemble
-from bsm.utils.general_utils import create_windowed_array
 from bsm.statistical_model.brnn_statistical_model import BRNNStatisticalModel
+from bsm.utils.general_utils import create_windowed_array
 from bsm.utils.normalization import Data
-from bsm.utils.type_aliases import StatisticalModelOutput
 
 key = jr.PRNGKey(0)
 input_dim = 1
@@ -61,7 +60,7 @@ def test_statistical_model_state_of_prediction():
 in_domain_test_ys = in_domain_test_ys.transpose()
 
 in_domain_preds = model.predict_batch(in_domain_test_xs,
-                                                                                      statistical_model_state)
+                                      statistical_model_state)
 
 
 def test_good_probabilistic_ensemble_fit():

tests/test_deterministic_gru_ensemble_bnn_sm.py

@@ -1,10 +1,8 @@
 import jax.numpy as jnp
 import jax.random as jr
-from jax import vmap
 
 from bsm.statistical_model import GPStatisticalModel
 from bsm.utils.normalization import Data
-from bsm.utils.type_aliases import StatisticalModelOutput
 
 key = jr.PRNGKey(0)
 input_dim = 1

tests/test_gaussian_processes_sm.py

@@ -1,11 +1,9 @@
 import jax.numpy as jnp
 import jax.random as jr
-from jax import vmap
 
 from bsm.bayesian_regression import ProbabilisticEnsemble
 from bsm.statistical_model import BNNStatisticalModel
 from bsm.utils.normalization import Data
-from bsm.utils.type_aliases import StatisticalModelOutput
 
 key = jr.PRNGKey(0)
 input_dim = 1

tests/test_probabilistic_ensemble_sm.py

@@ -1,11 +1,9 @@
 import jax.numpy as jnp
 import jax.random as jr
-from jax import vmap
 
 from bsm.bayesian_regression import ProbabilisticFSVGDEnsemble
 from bsm.statistical_model import BNNStatisticalModel
 from bsm.utils.normalization import Data
-from bsm.utils.type_aliases import StatisticalModelOutput
 
 key = jr.PRNGKey(0)
 input_dim = 1

tests/test_probabilistic_fsvg_ensemble_bnn_sm.py

@@ -6,7 +6,6 @@
 from bsm.statistical_model import BRNNStatisticalModel
 from bsm.utils.general_utils import create_windowed_array
 from bsm.utils.normalization import Data
-from bsm.utils.type_aliases import StatisticalModelOutput
 
 key = jr.PRNGKey(0)
 input_dim = 1