efax 2.2.4__tar.gz → 2.2.5__tar.gz

{efax-2.2.4 → efax-2.2.5}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: efax
-Version: 2.2.4
+Version: 2.2.5
 Summary: Exponential families for JAX
 Author: Neil Girdhar
 Author-email: Neil Girdhar <mistersheik@gmail.com>

{efax-2.2.4 → efax-2.2.5}/efax/_src/distributions/multivariate_normal/fixed_variance.py

@@ -76,7 +76,10 @@ class MultivariateFixedVarianceNormalNP(
     def sufficient_statistics(
         cls, x: JaxRealArray, **fixed_parameters: JaxArray
     ) -> MultivariateFixedVarianceNormalEP:
-        return MultivariateFixedVarianceNormalEP(x, variance=fixed_parameters["variance"])
+        xp = array_namespace(x)
+        return MultivariateFixedVarianceNormalEP(
+            x, variance=xp.broadcast_to(fixed_parameters["variance"], x.shape[:-1])
+        )
 
     @override
     def sample(self, key: KeyArray, shape: Shape | None = None) -> JaxRealArray:

{efax-2.2.4 → efax-2.2.5}/efax/_src/distributions/weibull.py

@@ -68,8 +68,8 @@ class WeibullNP(
     @classmethod
     def sufficient_statistics(cls, x: JaxRealArray, **fixed_parameters: JaxArray) -> WeibullEP:
         xp = array_namespace(x)
-        concentration = fixed_parameters["concentration"]
-        return WeibullEP(xp.broadcast_to(concentration, x.shape), x**concentration)
+        concentration = xp.broadcast_to(fixed_parameters["concentration"], x.shape)
+        return WeibullEP(concentration, x**concentration)
 
     @override
     def sample(self, key: KeyArray, shape: Shape | None = None) -> JaxRealArray:

{efax-2.2.4 → efax-2.2.5}/efax/_src/tools.py

@@ -38,13 +38,17 @@ T = TypeVar("T", bound=Distribution)
 
 
 def parameter_mean[T: Distribution](x: T, /, *, axis: Axis | None = None) -> T:
-    """Return the mean of the variable parameters; fixed parameters are preserved unchanged."""
+    """Return the mean over axis of all parameters, including fixed ones.
+
+    Fixed parameters are constant along the reduced axes (having been broadcast to the sample shape
+    in sufficient_statistics), so their mean equals their value.  Integer-typed parameters remain
+    integers after the mean because the Array API spec requires mean to preserve integer dtypes.
+    """
     xp = array_namespace(x)
     structure = Assembler.create_assembler(x)
-    fixed = parameters(x, fixed=True)
-    variable = parameters(x, fixed=False)
-    averaged = {path: xp.mean(value, axis=axis) for path, value in variable.items()}
-    return structure.assemble({**fixed, **averaged})
+    all_params = parameters(x)
+    averaged = {path: xp.mean(value, axis=axis) for path, value in all_params.items()}
+    return structure.assemble(averaged)
 
 
 def parameter_map[T: Distribution](

{efax-2.2.4 → efax-2.2.5}/pyproject.toml

@@ -17,7 +17,7 @@ dev = [
 
 [project]
 name = "efax"
-version = "2.2.4"
+version = "2.2.5"
 description = "Exponential families for JAX"
 readme = "README.rst"
 requires-python = ">=3.12, <3.15"