Only declare classes with fields as dataclasses

jax_transforms/alr.py

@@ -1,10 +1,7 @@
-from dataclasses import dataclass
-
 import jax
 import jax.numpy as jnp
 
 
-@dataclass
 class ALR:
     def unconstrain(self, x):
         return jnp.log(x[..., :-1]) - jnp.log(x[..., -1:])

jax_transforms/expanded_softmax.py

@@ -1,11 +1,8 @@
-from dataclasses import dataclass
-
 import jax
 import jax.numpy as jnp
 from tensorflow_probability.substrates.jax import distributions
 
 
-@dataclass
 class ExpandedSoftmax:
     def unconstrain(self, r_x):
         r, x = r_x

jax_transforms/ilr.py

@@ -1,5 +1,3 @@
-from dataclasses import dataclass, field
-
 import jax
 import jax.numpy as jnp
 
@@ -22,7 +20,6 @@ def _make_semiorthogonal_matrix(N: int):
     return V
 
 
-@dataclass
 class ILR:
     def unconstrain(self, x):
         N = x.shape[-1]

jax_transforms/normalized_gamma.py

@@ -41,7 +41,6 @@ def default_prior(self, x) -> distributions.Distribution:
         return distributions.ExpGamma(concentration=alpha_sum, rate=1)
 
 
-@dataclass
 class NormalizedExponential(NormalizedGamma):
     def __init__(self):
         super().__init__(alpha=jnp.ones(()))

jax_transforms/stickbreaking.py

@@ -7,7 +7,6 @@
 from .util import vmap_over_leading_axes
 
 
-@dataclass
 class StickbreakingBase:
     def _unconstrain_single(self, x):
         def _running_remainder(remainder, xi):
@@ -39,7 +38,6 @@ def constrain_with_logdetjac(self, z):
         return x, logJ
 
 
-@dataclass
 class StickbreakingCDF:
     def get_distribution(self, N: int) -> distributions.Distribution:
         raise NotImplementedError
@@ -65,13 +63,11 @@ def constrain_with_logdetjac(self, y):
         return x, logJ
 
 
-@dataclass
 class StickbreakingLogistic(StickbreakingCDF):
     def get_distribution(self, N: int) -> distributions.Logistic:
         return distributions.Logistic(loc=jnp.log(jnp.arange(N - 1, 0, -1)), scale=1)
 
 
-@dataclass
 class StickbreakingNormal(StickbreakingCDF):
     def get_distribution(self, N: int) -> distributions.Normal:
         return distributions.Normal(loc=jnp.log(jnp.arange(N - 1, 0, -1)) / 2, scale=1)
@@ -104,19 +100,16 @@ def constrain_with_logdetjac(self, y):
         return x, logJ
 
 
-@dataclass
 class StickbreakingPowerLogistic(StickbreakingPowerCDF):
     def __init__(self):
         super().__init__(distributions.Logistic(loc=0, scale=1))
 
 
-@dataclass
 class StickbreakingPowerNormal(StickbreakingPowerCDF):
     def __init__(self):
         super().__init__(distributions.Normal(loc=0, scale=1))
 
 
-@dataclass
 class StickbreakingAngular:
     def unconstrain(self, x):
         z = StickbreakingBase().unconstrain(x)