Fix imports and add calc_suff_stat/to_natural/to_canonical staticmethods.

Replace from-imports of torch.distributions classes with module-level
imports; use ch.distributions.* as base classes. Move calc_suff_stat,
to_natural, and to_canonical as @staticmethods onto each distribution
class. Remove redundant _calc_suff_stat overrides from truncated
distribution subclasses and inline direct calls where needed.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

Author: pstefanou12

delphi/distributions/boolean_product.py

@@ -2,15 +2,10 @@
 """Boolean product distribution in natural parameterization."""
 
 import torch as ch
-from torch.distributions import Bernoulli
+import torch.distributions
 
 
-def calc_bool_prod_suff_stat(x):
-    """Return sufficient statistics for boolean product distribution."""
-    return x
-
-
-class ExponentialFamilyBooleanProduct(Bernoulli):  # pylint: disable=abstract-method
+class ExponentialFamilyBooleanProduct(ch.distributions.Bernoulli):  # pylint: disable=abstract-method
     """Boolean product distribution parameterized by natural parameters."""
 
     def __init__(self, theta: ch.Tensor, dims: int):
@@ -19,6 +14,21 @@ def __init__(self, theta: ch.Tensor, dims: int):
         p = ch.exp(theta) / (1 + ch.exp(theta))
         super().__init__(p)
 
+    @staticmethod
+    def calc_suff_stat(x: ch.Tensor) -> ch.Tensor:
+        """Return sufficient statistics for boolean product distribution."""
+        return x
+
+    @staticmethod
+    def to_natural(theta: ch.Tensor) -> ch.Tensor:
+        """Convert canonical probability to natural log-odds parameter."""
+        return ch.log(theta / (1 - theta))
+
+    @staticmethod
+    def to_canonical(theta: ch.Tensor) -> ch.Tensor:
+        """Convert natural log-odds to canonical probability parameter."""
+        return ch.exp(theta) / (1 + ch.exp(theta))
+
     def log_prob(self, value):
         """Compute summed log probability over all dimensions."""
         result = super().log_prob(value)

delphi/distributions/exponential.py

@@ -2,15 +2,10 @@
 """Exponential distribution in natural parameterization."""
 
 import torch as ch
-from torch.distributions import Exponential
+import torch.distributions
 
 
-def calc_exp_suff_stat(x):
-    """Return sufficient statistics for exponential distribution."""
-    return x
-
-
-class ExponentialFamilyExponential(Exponential):  # pylint: disable=abstract-method
+class ExponentialFamilyExponential(ch.distributions.Exponential):  # pylint: disable=abstract-method
     """Exponential distribution parameterized by natural parameters."""
 
     def __init__(self, theta: ch.Tensor, dims: int):
@@ -19,6 +14,21 @@ def __init__(self, theta: ch.Tensor, dims: int):
         lambda_ = -theta
         super().__init__(lambda_)
 
+    @staticmethod
+    def calc_suff_stat(x: ch.Tensor) -> ch.Tensor:
+        """Return sufficient statistics for exponential distribution."""
+        return x
+
+    @staticmethod
+    def to_natural(theta: ch.Tensor) -> ch.Tensor:
+        """Convert canonical rate parameter to natural form."""
+        return -theta
+
+    @staticmethod
+    def to_canonical(theta: ch.Tensor) -> ch.Tensor:
+        """Convert natural parameters to canonical rate parameter."""
+        return -theta
+
     def log_prob(self, value):
         """Compute summed log probability over all dimensions."""
         result = super().log_prob(value)

delphi/distributions/multivariate_normal.py

@@ -2,45 +2,46 @@
 """Multivariate normal distribution in natural parameterization."""
 
 import torch as ch
-from torch import Tensor
-from torch.distributions import MultivariateNormal
+import torch.distributions
 
 
 class ExponentialFamilyMultivariateNormalKnownCovariance(  # pylint: disable=abstract-method
-    MultivariateNormal
+    ch.distributions.MultivariateNormal
 ):
     """Multivariate normal parameterized by natural parameters with known covariance."""
 
-    def __init__(self, covariance_matrix: Tensor, theta: Tensor, dims: int):
+    def __init__(self, covariance_matrix: ch.Tensor, theta: ch.Tensor, dims: int):
         """Initialize with covariance matrix, natural parameter theta, and dimension."""
         self.dims = dims
         v = theta
         mu = (covariance_matrix @ v).view(self.dims)
         super().__init__(mu, covariance_matrix)
 
     @staticmethod
-    def calc_suff_stat(x: Tensor) -> Tensor:
+    def calc_suff_stat(x: ch.Tensor) -> ch.Tensor:
         """Return sufficient statistics for multivariate normal with known covariance."""
         return x
 
     @staticmethod
-    def to_natural(theta: Tensor, covariance_matrix: Tensor) -> Tensor:
+    def to_natural(theta: ch.Tensor, covariance_matrix: ch.Tensor) -> ch.Tensor:
         """Convert canonical mean to natural parameters."""
         inv_cov = covariance_matrix.inverse()
         v = theta @ inv_cov  # pylint: disable=invalid-name
         return v.flatten()
 
     @staticmethod
-    def to_canonical(theta: Tensor, covariance_matrix: Tensor) -> Tensor:
+    def to_canonical(theta: ch.Tensor, covariance_matrix: ch.Tensor) -> ch.Tensor:
         """Convert natural parameters to canonical mean."""
         loc = theta @ covariance_matrix
         return loc.flatten()
 
 
-class ExponentialFamilyMultivariateNormal(MultivariateNormal):  # pylint: disable=abstract-method
+class ExponentialFamilyMultivariateNormal(  # pylint: disable=abstract-method
+    ch.distributions.MultivariateNormal
+):
     """Multivariate normal parameterized by natural parameters."""
 
-    def __init__(self, theta: Tensor, dims: int):
+    def __init__(self, theta: ch.Tensor, dims: int):
         """Initialize with natural parameter theta and dimension."""
         self.dims = dims
         T, v = theta[: self.dims**2], theta[self.dims**2 :]  # pylint: disable=invalid-name
@@ -49,12 +50,12 @@ def __init__(self, theta: Tensor, dims: int):
         super().__init__(mu, covariance_matrix)
 
     @staticmethod
-    def calc_suff_stat(x: Tensor) -> Tensor:
+    def calc_suff_stat(x: ch.Tensor) -> ch.Tensor:
         """Return sufficient statistics for multivariate normal."""
         return ch.cat([ch.bmm(x.unsqueeze(2), x.unsqueeze(1)).flatten(1), x], 1)
 
     @staticmethod
-    def to_natural(theta: Tensor, dims: int) -> Tensor:
+    def to_natural(theta: ch.Tensor, dims: int) -> ch.Tensor:
         """Convert canonical parameters to natural form."""
         cov_matrix = theta[: dims**2].view(dims, dims)
         loc = theta[dims**2 :]
@@ -63,7 +64,7 @@ def to_natural(theta: Tensor, dims: int) -> Tensor:
         return ch.cat([-0.5 * mat_t.flatten(), v.flatten()])
 
     @staticmethod
-    def to_canonical(theta: Tensor, dims: int) -> Tensor:
+    def to_canonical(theta: ch.Tensor, dims: int) -> ch.Tensor:
         """Convert natural parameters to canonical form."""
         mat_t = theta[: dims**2].view(dims, dims)  # pylint: disable=invalid-name
         v = theta[dims**2 :]

delphi/distributions/poisson.py

@@ -2,15 +2,10 @@
 """Poisson distribution in natural parameterization."""
 
 import torch as ch
-from torch.distributions import Poisson
+import torch.distributions
 
 
-def calc_poiss_suff_stat(x):
-    """Return sufficient statistics for Poisson distribution."""
-    return x
-
-
-class ExponentialFamilyPoisson(Poisson):  # pylint: disable=abstract-method
+class ExponentialFamilyPoisson(ch.distributions.Poisson):  # pylint: disable=abstract-method
     """Poisson distribution parameterized by natural parameters."""
 
     def __init__(self, theta: ch.Tensor, dims: int):
@@ -19,6 +14,21 @@ def __init__(self, theta: ch.Tensor, dims: int):
         lambda_ = ch.exp(theta)
         super().__init__(lambda_)
 
+    @staticmethod
+    def calc_suff_stat(x: ch.Tensor) -> ch.Tensor:
+        """Return sufficient statistics for Poisson distribution."""
+        return x
+
+    @staticmethod
+    def to_natural(theta: ch.Tensor) -> ch.Tensor:
+        """Convert canonical rate parameter to natural log form."""
+        return ch.log(theta)
+
+    @staticmethod
+    def to_canonical(theta: ch.Tensor) -> ch.Tensor:
+        """Convert natural parameters to canonical rate parameter."""
+        return ch.exp(theta)
+
     def log_prob(self, value):
         """Compute summed log probability over all dimensions."""
         result = super().log_prob(value)

delphi/distributions/weibull.py

@@ -2,15 +2,10 @@
 """Weibull distribution in natural parameterization."""
 
 import torch as ch
-from torch.distributions import Weibull
+import torch.distributions
 
 
-def calc_weibull_suff_stat(k, x):
-    """Return sufficient statistics for Weibull distribution."""
-    return x.pow(k)
-
-
-class ExponentialFamilyWeibull(Weibull):  # pylint: disable=abstract-method
+class ExponentialFamilyWeibull(ch.distributions.Weibull):  # pylint: disable=abstract-method
     """Weibull distribution parameterized by natural parameters."""
 
     def __init__(self, k: ch.Tensor, theta: ch.Tensor, dims: int):
@@ -19,6 +14,21 @@ def __init__(self, k: ch.Tensor, theta: ch.Tensor, dims: int):
         lambda_ = (-1 / theta).pow(1 / k)
         super().__init__(lambda_, k)
 
+    @staticmethod
+    def calc_suff_stat(k: ch.Tensor, x: ch.Tensor) -> ch.Tensor:
+        """Return sufficient statistics for Weibull distribution."""
+        return x.pow(k)
+
+    @staticmethod
+    def to_natural(k: ch.Tensor, theta: ch.Tensor) -> ch.Tensor:
+        """Convert canonical scale parameter to natural form."""
+        return -1.0 / theta.pow(k)
+
+    @staticmethod
+    def to_canonical(k: ch.Tensor, theta: ch.Tensor) -> ch.Tensor:
+        """Convert natural parameters to canonical scale parameter."""
+        return (-1 / theta).pow(1 / k)
+
     def log_prob(self, value):
         """Compute summed log probability over all dimensions."""
         result = super().log_prob(value)

delphi/truncated/distributions/truncated_boolean_product.py

@@ -49,11 +49,6 @@ def __init__(
             logger,
         )
 
-    @staticmethod
-    def _calc_suff_stat(x):
-        """Compute sufficient statistics for the Boolean product distribution."""
-        return boolean_product.calc_bool_prod_suff_stat(x)
-
     def _reparameterize_nat_form(self, theta):
         """Convert canonical probability parameter to natural log-odds form."""
         return ch.log(theta / (1 - theta))

delphi/truncated/distributions/truncated_exponential.py

@@ -53,11 +53,6 @@ def __init__(
             logger,
         )
 
-    @staticmethod
-    def _calc_suff_stat(x):
-        """Compute sufficient statistics for the exponential distribution."""
-        return exponential.calc_exp_suff_stat(x)
-
     def _constraints(self, theta):
         """Clamp theta to be strictly negative."""
         return ch.clamp(theta, max=-1e-6)

delphi/truncated/distributions/truncated_multivariate_normal.py

@@ -65,15 +65,14 @@ def __init__(
         self.emp_T = None  # pylint: disable=invalid-name
         self.emp_v = None
 
-    @staticmethod
-    def _calc_suff_stat(x: ch.Tensor) -> ch.Tensor:
-        """Compute sufficient statistics for multivariate normal."""
-        return multivariate_normal.ExponentialFamilyMultivariateNormal.calc_suff_stat(x)
-
     def _calc_emp_model(self):
         """Calculate empirical natural parameters and register T and v as nn.Parameters."""
         dataset_s = self.train_loader_.dataset.S  # pylint: disable=invalid-name
-        suff_stats = self._calc_suff_stat(dataset_s).mean(0)
+        suff_stats = (
+            multivariate_normal.ExponentialFamilyMultivariateNormal.calc_suff_stat(
+                dataset_s
+            ).mean(0)
+        )
         second_moment = suff_stats[: self.dims**2].view(self.dims, self.dims)
         loc = suff_stats[self.dims**2 :]
         # Center the second moment to get the empirical covariance: Σ = E[xx^T] - μμ^T.

delphi/truncated/distributions/truncated_multivariate_normal_known_covariance.py

@@ -62,17 +62,12 @@ def __init__(
         self.covariance_matrix = covariance_matrix
         self._sampler = sampler
 
-    @staticmethod
-    def _calc_suff_stat(x: ch.Tensor) -> ch.Tensor:
-        """Compute sufficient statistics for multivariate normal with known covariance."""
-        return multivariate_normal.ExponentialFamilyMultivariateNormalKnownCovariance.calc_suff_stat(
-            x
-        )
-
     def _calc_emp_model(self):
         """Calculate empirical natural parameters and register theta as an nn.Parameter."""
         dataset_s = self.train_loader_.dataset.S  # pylint: disable=invalid-name
-        emp_mean = self._calc_suff_stat(dataset_s).mean(0)
+        emp_mean = multivariate_normal.ExponentialFamilyMultivariateNormalKnownCovariance.calc_suff_stat(
+            dataset_s
+        ).mean(0)
         self.emp_theta = multivariate_normal.ExponentialFamilyMultivariateNormalKnownCovariance.to_natural(
             emp_mean, self.covariance_matrix
         )

delphi/truncated/distributions/truncated_poisson.py

@@ -53,11 +53,6 @@ def __init__(
             logger,
         )
 
-    @staticmethod
-    def _calc_suff_stat(x):
-        """Compute sufficient statistics for the Poisson distribution."""
-        return poisson.calc_poiss_suff_stat(x)
-
     def _reparameterize_nat_form(self, theta):
         """Convert canonical rate parameter to natural log form."""
         return ch.log(theta)