exact_log_prob keyword, rademacher noise, examples install packages, version

Author: homerjed

README.md

@@ -77,17 +77,21 @@ Install via
 pip install sbgm
 ```
 
-See [examples](https://github.com/homerjed/sbgm/tree/main/examples).
+and for the [examples](https://github.com/homerjed/sbgm/tree/main/examples), run 
 
-To run on the `cifar10` image dataset, try something like
+```
+pip install .[examples] 
+```
+
+To fit a diffusion model to the `cifar10` image dataset, try something like
 
 ```python
 import sbgm
 import data
 import configs
 
-datasets_path = "."
-root_dir = "."
+datasets_path = "./"
+root_dir = "./"
 
 config = configs.cifar10_config()
 
@@ -129,7 +133,7 @@ model = sbgm.train.train(
 * UNet and transformer score network implementations,
 * VP, SubVP and VE SDEs (neural network $\beta(t)$ and $\sigma(t)$ functions are on the list!),
 * Multi-modal conditioning (basically just optional parameter and image conditioning methods),
-* Checkpointing optimiser and model,
+* Checkpointing for optimiser and model,
 * Multi-device training and sampling.
 
 ### Samples

configs/cifar10.py

@@ -35,7 +35,7 @@ def cifar10_config():
     # Sampling
     config.use_ema               = False
     config.sample_size           = 5
-    config.exact_logp            = False
+    config.exact_log_prob            = False
     config.ode_sample            = True
     config.eu_sample             = True
 

configs/flowers.py

@@ -35,7 +35,7 @@ def flowers_config():
 
     # Sampling
     config.sample_size           = 5
-    config.exact_logp            = False
+    config.exact_log_prob            = False
     config.ode_sample            = True
     config.eu_sample             = True
 

configs/grfs.py

@@ -34,7 +34,7 @@ def grfs_config():
     # Sampling
     config.use_ema         = False
     config.sample_size     = 5
-    config.exact_logp      = False
+    config.exact_log_prob      = False
     config.ode_sample      = True
     config.eu_sample       = True
 

configs/mnist.py

@@ -34,7 +34,7 @@ def mnist_config():
 
     # Sampling
     config.sample_size           = 8
-    config.exact_logp            = False
+    config.exact_log_prob        = False
     config.ode_sample            = True
     config.eu_sample             = True
     config.use_ema               = False

configs/moons.py

@@ -31,7 +31,7 @@ def moons_config():
     # Sampling
     config.use_ema               = True
     config.sample_size           = 64 # Squared in sampling
-    config.exact_logp            = True
+    config.exact_log_prob        = True
     config.ode_sample            = False 
     config.eu_sample             = False 
 

configs/quijote.py

@@ -34,7 +34,7 @@ def quijote_config():
     # Sampling
     config.use_ema               = False
     config.sample_size           = 5
-    config.exact_logp            = False
+    config.exact_log_prob        = False
     config.ode_sample            = True
     config.eu_sample             = True
 

examples/cifar10.ipynb

@@ -36,7 +36,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 3,
+   "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -78,7 +78,7 @@
     "# Sampling\n",
     "use_ema               = False\n",
     "sample_size           = 5 # Squared for a grid\n",
-    "exact_logp            = False\n",
+    "exact_log_prob        = False\n",
     "ode_sample            = True # Sample the ODE during training\n",
     "eu_sample             = True # Euler-Maruyama sample the SDE during training\n",
     "\n",
@@ -459,7 +459,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 13,
+   "execution_count": null,
    "metadata": {},
    "outputs": [
     {
@@ -475,7 +475,7 @@
    ],
    "source": [
     "log_likelihood_fn = sbgm.ode.get_log_likelihood_fn(\n",
-    "    model, sde, dataset.data_shape, exact_logp=False, n_eps=64\n",
+    "    model, sde, data_shape=dataset.data_shape, exact_log_prob=False, n_eps=64\n",
     ")\n",
     "L_X = log_likelihood_fn(X[0], None, A[0], key)\n",
     "\n",

examples/grfs.ipynb

@@ -44,7 +44,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 3,
+   "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -82,7 +82,7 @@
     "# Sampling\n",
     "use_ema               = False\n",
     "sample_size           = 5 # Squared for a grid\n",
-    "exact_logp            = False\n",
+    "exact_log_prob            = False\n",
     "ode_sample            = True # Sample the ODE during training\n",
     "eu_sample             = True # Euler-Maruyama sample the SDE during training\n",
     "\n",
@@ -415,7 +415,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 11,
+   "execution_count": null,
    "metadata": {},
    "outputs": [
     {
@@ -433,7 +433,7 @@
     "key, key_L = jr.split(key)\n",
     "\n",
     "log_likelihood_fn = sbgm.ode.get_log_likelihood_fn(\n",
-    "    model, sde, dataset.data_shape, exact_logp=True\n",
+    "    model, sde, data_shape=dataset.data_shape, exact_log_prob=True\n",
     ")\n",
     "L_X = log_likelihood_fn(X[0], Q[0], A[0], key_L)\n",
     "\n",

examples/grfs.py

@@ -56,7 +56,7 @@
 # Sampling
 use_ema               = False
 sample_size           = 5 # Squared for a grid
-exact_logp            = False
+exact_log_prob            = False
 ode_sample            = True # Sample the ODE during training
 eu_sample             = True # Euler-Maruyama sample the SDE during training
 
@@ -245,7 +245,7 @@ def diffuse(x, t, eps):
 key, key_L = jr.split(key)
 
 log_likelihood_fn = sbgm.ode.get_log_likelihood_fn(
-    model, sde, dataset.data_shape, exact_logp=True
+    model, sde, dataset.data_shape, exact_log_prob=True
 )
 L_X = log_likelihood_fn(X[0], Q[0], A[0], key_L)