init

.gitignore

@@ -0,0 +1,130 @@
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+*.ipynb
+notebooks/
+# C extensions
+*.so
+data/
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+pip-wheel-metadata/
+share/python-wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# PyInstaller
+#  Usually these files are written by a python script from a template
+#  before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.nox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+*.py,cover
+.hypothesis/
+.pytest_cache/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+db.sqlite3
+db.sqlite3-journal
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# IPython
+profile_default/
+ipython_config.py
+
+# pyenv
+.python-version
+
+# pipenv
+#   According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
+#   However, in case of collaboration, if having platform-specific dependencies or dependencies
+#   having no cross-platform support, pipenv may install dependencies that don't work, or not
+#   install all needed dependencies.
+#Pipfile.lock
+
+# PEP 582; used by e.g. github.com/David-OConnor/pyflow
+__pypackages__/
+
+# Celery stuff
+celerybeat-schedule
+celerybeat.pid
+
+# SageMath parsed files
+*.sage.py
+
+# Environments
+.env
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+
+# Spyder project settings
+.spyderproject
+.spyproject
+
+# Rope project settings
+.ropeproject
+
+# mkdocs documentation
+/site
+
+# mypy
+.mypy_cache/
+.dmypy.json
+dmypy.json
+
+# Pyre type checker
+.pyre/

LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2021 James Thornton
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

README.md

@@ -0,0 +1 @@
+# diffusions

diffusions/datasets.py

@@ -0,0 +1,6 @@
+def central_scalar(x):
+    return lambda x: x * 2. - 1.
+
+def inverse_central_scalar(x):
+    return lambda x: (x + 1.) / 2.
+

diffusions/diffusion_spec/base.py

@@ -0,0 +1,25 @@
+
+import jax
+from jax import numpy as jnp
+from jax import random
+from ..model_ioputs import DiffusionOutput, DiffusionModelOutput
+from ..typing import KeyArray, Function
+
+class ReverseDiffusionSpec:
+
+    def sample_t(self, rng: KeyArray, x_0: jnp.ndarray) -> jnp.ndarray:
+        pass
+
+    def sample_x(self, x_0 : jnp.ndarray, t: jnp.ndarray) -> jnp.ndarray:
+        pass
+
+    def forward_sample(self, rng: KeyArray, x_0: jnp.ndarray) -> DiffusionOutput:
+        rng, t_rng, x_rng = random.split(rng, 3)
+        t = self.sample_t(t_rng, x_0)
+        return self.sample_x(x_rng, x_0, t)
+
+    def pointwise_loss(self, model_output: DiffusionModelOutput, perturbed_output: DiffusionOutput) -> jnp.ndarray:
+        pass
+
+    def simulate_reverse_diffusion(self, rng: KeyArray, x_T: jnp.ndarray, score_fn: Function):
+        pass

diffusions/diffusion_spec/discrete_ou.py

@@ -0,0 +1,71 @@
+import imp
+import jax
+from jax import numpy as jnp
+from jax import random
+from jtils.batch_ops import batch_mul
+from ..model_ioputs import DiffusionOutput, DiffusionModelOutput, DiffusionModelInput
+from .base import ReverseDiffusionSpec
+from .losses import mean_squared_score_loss
+from ..typing import KeyArray, Function
+from .simulate_diffusion import get_mean_scale_reverse_fn
+
+
+class DiscreteOU(ReverseDiffusionSpec):
+    """ """
+
+    def __init__(self, beta_min=0.1, beta_max=20, N=1000, eps=1e-3, T=1.0) -> None:
+        super().__init__()
+        self.beta_0 = beta_min
+        self.beta_1 = beta_max
+        self.N = N
+        self.T = T
+        self.eps = eps
+
+        self.discrete_betas = jnp.linspace(beta_min / N, beta_max / N, N)
+        self.alphas = 1.0 - self.discrete_betas
+        self.alphas_cumprod = jnp.cumprod(self.alphas, axis=0)
+        self.sqrt_alphas_cumprod = jnp.sqrt(self.alphas_cumprod)
+        self.sqrt_1m_alphas_cumprod = jnp.sqrt(1.0 - self.alphas_cumprod)
+
+    def sample_t(self, rng: KeyArray, x_0: jnp.ndarray) -> jnp.ndarray:
+        return jax.random.choice(rng, self.N, (x_0.shape[0],))
+
+    def mean_scale_function(
+        self, rng: KeyArray, x_t: jnp.ndarray, t: jnp.ndarray, score_fn: Function
+    ) -> jnp.ndarray:
+        timestep = t * (self.N - 1) / self.T
+        t_label = timestep.astype(jnp.int32)
+        beta = self.discrete_betas[t_label]
+        inputs = DiffusionModelInput(x_t=x_t, t=timestep)
+        model_pred = score_fn(inputs)
+        model_pred = model_pred.target
+        std = self.sqrt_1m_alphas_cumprod[t_label.astype(jnp.int32)]
+        score = batch_mul(-model_pred, 1.0 / std)
+        x_mean = batch_mul((x_t + batch_mul(beta, score)), 1.0 / jnp.sqrt(1.0 - beta))
+        return x_mean, jnp.sqrt(beta)
+
+    def sample_x(self, rng: KeyArray, x_0: jnp.ndarray, t: jnp.ndarray) -> jnp.ndarray:
+        noise = random.normal(rng, x_0.shape)
+        x_t = batch_mul(self.sqrt_alphas_cumprod[t], x_0) + batch_mul(
+            self.sqrt_1m_alphas_cumprod[t], noise
+        )
+        return DiffusionOutput(x_t=x_t, z=noise, t=t)
+
+    def pointwise_loss(
+        self, model_output: DiffusionModelOutput, perturbed_output: DiffusionOutput
+    ) -> jnp.ndarray:
+        return mean_squared_score_loss(
+            model_output=model_output, perturbed_output=perturbed_output
+        )
+
+    def get_simulate_reverse_diffusion_fn(self, score_fn: Function):
+        return get_mean_scale_reverse_fn(
+            score_fn, self.mean_scale_function, self.eps, self.T, self.N
+        )
+
+    def simulate_reverse_diffusion(
+        self, rng: KeyArray, x_T: jnp.ndarray, score_fn: Function
+    ):
+        return get_mean_scale_reverse_fn(
+            score_fn, self.mean_scale_function, self.eps, self.T, self.N
+        )(rng, x_T)

diffusions/diffusion_spec/losses.py

@@ -0,0 +1,16 @@
+from jax import numpy as jnp
+from ..model_ioputs import DiffusionOutput, DiffusionModelOutput
+
+
+def mean_squared_score_loss(
+    model_output: DiffusionModelOutput, perturbed_output: DiffusionOutput
+) -> jnp.ndarray:
+
+    noise = perturbed_output.z
+    model_pred = model_output.target
+
+    losses = jnp.square(model_pred - noise)
+    losses = jnp.mean(losses.reshape((losses.shape[0], -1)), axis=-1)
+    loss = jnp.mean(losses)
+
+    return loss

diffusions/diffusion_spec/simulate_diffusion.py

@@ -0,0 +1,38 @@
+
+
+import jax
+from jax import random
+from jax import numpy as jnp
+from jtils.batch_ops import batch_add, batch_mul
+from ..typing import Function, KeyArray
+
+def get_mean_scale_reverse_fn(score_fn: Function, 
+                              mean_scale_fn: Function, 
+                              eps: float = 1e-3, 
+                              T: float=1.0, 
+                              N: int=1000):
+
+    def simulate_reverse_diffusion(rng: KeyArray, x_T: jnp.ndarray):
+        shape = x_T.shape
+        def update_fn(rng, score_fn, x_t, t):
+            rng, step_rng= random.split(rng)
+            x_mean, scale = mean_scale_fn(step_rng, x_t, t, score_fn)
+            noise = random.normal(rng, x_t.shape)
+            x = x_mean + batch_mul(scale, noise)
+            return x, x_mean
+
+        timesteps = jnp.linspace(T, eps, N)
+
+        def loop_body(i, val):
+            rng, x, x_mean = val
+            t = timesteps[i]
+            vec_t = jnp.ones(shape[0]) * t
+            rng, step_rng = random.split(rng)
+            x, x_mean = update_fn(step_rng, score_fn, x, vec_t)
+            return rng, x, x_mean
+
+        _, x, x_mean = jax.lax.fori_loop(0, N, loop_body, (rng, x_T, x_T))
+
+        return x, x_mean
+
+    return simulate_reverse_diffusion

diffusions/model_ioputs.py

@@ -0,0 +1,60 @@
+import jax
+from typing import Optional, Tuple
+import flax
+import flax.linen as nn
+import jax
+import jax.numpy as jnp
+from flax.core.frozen_dict import FrozenDict, unfreeze
+from collections import OrderedDict
+from .typing import KeyArray
+from jtils.namespace import dict_to_namespace
+from argparse import Namespace
+from jtils.namespace import namespace_to_dict
+
+
+class DataClass(OrderedDict):
+    def __getitem__(self, key):
+        return self.__getattribute__(key)
+
+
+@flax.struct.dataclass
+class DiffusionOutput(DataClass):
+    """ """
+
+    x_t: jnp.ndarray = None
+    x_0: Optional[jnp.ndarray] = None
+    z: Optional[jnp.ndarray] = None
+    t: Optional[jnp.ndarray] = None
+    latent: Optional[jnp.ndarray] = None
+    label: Optional[jnp.ndarray] = None
+
+
+@flax.struct.dataclass
+class DiffusionModelOutput(DataClass):
+    """ """
+
+    target: jnp.ndarray = None
+    x_next: Optional[jnp.ndarray] = None
+    x_0: Optional[jnp.ndarray] = None
+
+
+@flax.struct.dataclass
+class DiffusionModelInput(DataClass):
+    """ """
+
+    x_t: jnp.ndarray = None
+    t: jnp.ndarray = None
+    rng : Optional[KeyArray] = None
+    x_0: Optional[jnp.ndarray] = None
+
+
+
+class ModelConfig(Namespace):
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
+
+    def items(self):
+        return namespace_to_dict(self).items()
+
+    def __getitem__(self, key):
+        return self.__getattribute__(key)

diffusions/model_utils.py

@@ -0,0 +1,14 @@
+from typing import Callable
+import jax
+from jax import numpy as jnp
+from flax import linen as nn
+from .typing import Params, Function
+from .model_ioputs import DiffusionModelOutput
+from .models.base import DiffusionModel
+
+def get_model_fn(model: DiffusionModel, params: Params) -> Function:
+    def model_fn(DiffusionModelInput) -> DiffusionModelOutput:
+        model_output = model.apply({"params": params}, DiffusionModelInput)
+        return model_output
+
+    return model_fn

diffusions/models/__init__.py

@@ -0,0 +1 @@
+from .mlp import MLPNet