WIP: Compact interval fixes

bgflow/distribution/distributions.py

@@ -2,9 +2,10 @@
 import torch
 from .energy import Energy
 from .sampling import Sampler
+from torch.distributions import constraints
 
 
-__all__ = ["TorchDistribution", "CustomDistribution", "UniformDistribution"]
+__all__ = ["TorchDistribution", "CustomDistribution", "UniformDistribution", "SloppyUniform"]
 
 
 class CustomDistribution(Energy, Sampler):
@@ -67,9 +68,47 @@ def __getattr__(self, name):
             raise AttributeError(msg)
 
 
+class _SloppyUniform(torch.distributions.Uniform):
+    def __init__(self, *args, tol=1e-5, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.tol = tol
+
+    @constraints.dependent_property(is_discrete=False, event_dim=0)
+    def support(self):
+        return constraints.interval(self.low-self.tol, self.high+self.tol)
+
+
+class SloppyUniform(torch.nn.Module):
+    def __init__(self, low, high, validate_args=None, tol=1e-5):
+        super().__init__()
+        self.register_buffer("low", low)
+        self.register_buffer("high", high)
+        self.tol = tol
+        self.validate_args = validate_args
+
+    def __getattr__(self, name):
+        try:
+            return super().__getattr__(name=name)
+        except AttributeError:
+            uniform = _SloppyUniform(self.low, self.high, self.validate_args, tol=self.tol)
+            if hasattr(uniform, name):
+                return getattr(uniform, name)
+        except:
+            raise AttributeError(f"SloppyUniform has no attribute {name}")
+
+
 class UniformDistribution(TorchDistribution):
     """Shortcut"""
-    def __init__(self, low, high, validate_args=None, n_event_dims=1):
-        uniform = torch.distributions.Uniform(low, high, validate_args)
+    def __init__(self, low, high, tol=1e-5, validate_args=None, n_event_dims=1):
+        uniform = SloppyUniform(low, high, validate_args, tol=tol)
         independent = torch.distributions.Independent(uniform, n_event_dims)
-        super().__init__(independent)
+        super().__init__(independent)
+        self.uniform = uniform
+
+    def _energy(self, x):
+        try:
+            y = - self._delegate.log_prob(x)[:,None]
+            assert torch.all(torch.isfinite(y))
+            return y
+        except (ValueError, AssertionError):
+            return -self._delegate.log_prob(self._delegate.sample(sample_shape=x.shape[:-1]))[:,None]

bgflow/distribution/modulify.py

@@ -0,0 +1,60 @@
+import torch
+from torch import nn
+from torch.distributions.utils import lazy_property
+
+
+__all__ = ["distribution_module"]
+
+
+def distribution_module(cls):
+    """A class decorator to "modulify" torch distributions."""
+    # get the actual distribution class from the mro()
+    cls_distribution = []
+    for parent in cls.mro()[1:]:
+        if (
+            issubclass(parent, torch.distributions.Distribution)
+            and parent != torch.distributions.Distribution
+        ):
+            cls_distribution.append(parent)
+    assert len(cls_distribution) == 1
+    cls_distribution = cls_distribution[0]
+
+    class DistributionModule:
+        def __init__(self, *args, **kwargs):
+            """Initialize distribution and register plain tensors as buffers."""
+            super().__init__(*args, **kwargs)
+
+            # after initializing re-register distribution parameters as buffers
+            k_tensors = []
+            for k in self.__dict__:
+
+                # check if it's a lazy_property
+                # in this case we should simply move on instead of evaluating it.
+                if self._is_lazy_property(k):
+                    continue
+
+                # check if it's a "plain" tensor
+                if isinstance(getattr(self, k), torch.Tensor):
+                    k_tensors.append(k)
+
+            for k in k_tensors:
+                # for a "plain "tensor we will now register it as buffer.
+                val = getattr(self, k)
+                delattr(self, k)
+                self.register_buffer(k, val)
+
+        def __getattribute__(self, name):
+            """Return attribute with lazy_property special check."""
+            if type(self)._is_lazy_property(name):
+                # deleting the attribute from the instance will simply "reset"
+                # the lazy_property
+                delattr(self, name)
+            return super().__getattribute__(name)
+
+        @classmethod
+        def _is_lazy_property(cls, name):
+            return isinstance(getattr(cls, name, object), lazy_property)
+
+    return type(
+        cls.__name__, (DistributionModule, cls_distribution, nn.Module), {}
+    )

bgflow/nn/dense.py

@@ -4,7 +4,7 @@
 from ..utils.types import is_list_or_tuple
 
 
-__all__ = ["DenseNet", "MeanFreeDenseNet"]
+__all__ = ["DenseNet", "MeanFreeDenseNet", "SirenDenseNet"]
 
 
 class DenseNet(torch.nn.Module):
@@ -53,3 +53,28 @@ class MeanFreeDenseNet(DenseNet):
     def forward(self, x):
         y = self._layers(x)
         return y - y.mean(dim=1, keepdim=True)
+
+
+class Sin(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+
+    def forward(self, x):
+        return torch.sin(x)
+
+
+class SirenDenseNet(DenseNet):
+    def __init__(self, *args, scale_first_weights=True, initialize=True, **kwargs):
+        super().__init__(*args, **kwargs, activation=Sin())
+        if initialize:
+            self._init_siren_weights(self._layers, scale_first_weights)
+
+    @staticmethod
+    def _init_siren_weights(layers, scale_first_weights=True):
+        with torch.no_grad():
+            linear_layers = [layer for layer in layers if isinstance(layer, torch.nn.Linear)]
+            for layer in linear_layers:
+                n = layer.weight.shape[-1]
+                layer.weight.data = -np.sqrt(6./n) + 2.0*np.sqrt(6./n) * torch.rand_like(layer.weight.data)
+            if scale_first_weights:
+                linear_layers[0].weight.data *= 30.

bgflow/nn/flow/__init__.py

@@ -4,6 +4,7 @@
 from .estimator import *
 from .stochastic import *
 from .transformer import *
+from .root_finding import *
 
 from .affine import *
 from .coupling import *

bgflow/nn/flow/root_finding/__init__.py

@@ -0,0 +1,3 @@
+from .approx_inverse import *
+from .bisection import *
+from .newton import *