aspire-inference 0.1.0a2__tar.gz → 0.1.0a4__tar.gz

{aspire_inference-0.1.0a2/aspire_inference.egg-info → aspire_inference-0.1.0a4}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: aspire-inference
-Version: 0.1.0a2
+Version: 0.1.0a4
 Summary: Accelerate Sequential Posterior Inference via REuse
 Author-email: "Michael J. Williams" <michaeljw1@googlemail.com>
 License: MIT
@@ -14,6 +14,8 @@ Requires-Dist: numpy
 Requires-Dist: array-api-compat
 Requires-Dist: wrapt
 Requires-Dist: h5py
+Provides-Extra: scipy
+Requires-Dist: scipy; extra == "scipy"
 Provides-Extra: jax
 Requires-Dist: jax; extra == "jax"
 Requires-Dist: jaxlib; extra == "jax"
@@ -21,6 +23,7 @@ Requires-Dist: flowjax; extra == "jax"
 Provides-Extra: torch
 Requires-Dist: torch; extra == "torch"
 Requires-Dist: zuko; extra == "torch"
+Requires-Dist: tqdm; extra == "torch"
 Provides-Extra: minipcn
 Requires-Dist: minipcn; extra == "minipcn"
 Provides-Extra: emcee

{aspire_inference-0.1.0a2 → aspire_inference-0.1.0a4/aspire_inference.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: aspire-inference
-Version: 0.1.0a2
+Version: 0.1.0a4
 Summary: Accelerate Sequential Posterior Inference via REuse
 Author-email: "Michael J. Williams" <michaeljw1@googlemail.com>
 License: MIT
@@ -14,6 +14,8 @@ Requires-Dist: numpy
 Requires-Dist: array-api-compat
 Requires-Dist: wrapt
 Requires-Dist: h5py
+Provides-Extra: scipy
+Requires-Dist: scipy; extra == "scipy"
 Provides-Extra: jax
 Requires-Dist: jax; extra == "jax"
 Requires-Dist: jaxlib; extra == "jax"
@@ -21,6 +23,7 @@ Requires-Dist: flowjax; extra == "jax"
 Provides-Extra: torch
 Requires-Dist: torch; extra == "torch"
 Requires-Dist: zuko; extra == "torch"
+Requires-Dist: tqdm; extra == "torch"
 Provides-Extra: minipcn
 Requires-Dist: minipcn; extra == "minipcn"
 Provides-Extra: emcee

{aspire_inference-0.1.0a2 → aspire_inference-0.1.0a4}/aspire_inference.egg-info/requires.txt

@@ -18,6 +18,9 @@ flowjax
 [minipcn]
 minipcn
 
+[scipy]
+scipy
+
 [test]
 pytest
 pytest-requires
@@ -25,3 +28,4 @@ pytest-requires
 [torch]
 torch
 zuko
+tqdm

{aspire_inference-0.1.0a2 → aspire_inference-0.1.0a4}/pyproject.toml

@@ -25,6 +25,9 @@ dependencies = [
 dynamic = ["version"]
 
 [project.optional-dependencies]
+scipy = [
+    "scipy",
+]
 jax = [
     "jax",
     "jaxlib",
@@ -33,6 +36,7 @@ jax = [
 torch = [
     "torch",
     "zuko",
+    "tqdm",
 ]
 minipcn = [
     "minipcn",

{aspire_inference-0.1.0a2 → aspire_inference-0.1.0a4}/src/aspire/flows/base.py

@@ -1,18 +1,28 @@
+import logging
 from typing import Any
 
 from ..history import FlowHistory
-from ..transforms import BaseTransform
+from ..transforms import BaseTransform, IdentityTransform
+
+logger = logging.getLogger(__name__)
 
 
 class Flow:
+    xp = None  # type: Any
+
     def __init__(
         self,
         dims: int,
         device: Any,
-        data_transform: BaseTransform = None,
+        data_transform: BaseTransform | None = None,
     ):
         self.dims = dims
         self.device = device
+
+        if data_transform is None:
+            data_transform = IdentityTransform(self.xp)
+            logger.info("No data_transform provided, using IdentityTransform.")
+
         self.data_transform = data_transform
 
     def log_prob(self, x):

{aspire_inference-0.1.0a2 → aspire_inference-0.1.0a4}/src/aspire/transforms.py

@@ -4,7 +4,6 @@ from typing import Any
 
 from array_api_compat import device as get_device
 from array_api_compat import is_torch_namespace
-from scipy.special import erf, erfinv
 
 from .flows import get_flow_wrapper
 from .utils import (
@@ -50,13 +49,17 @@ class IdentityTransform(BaseTransform):
     """Identity transform that does nothing to the data."""
 
     def fit(self, x):
-        return x
+        return copy_array(x, xp=self.xp)
 
     def forward(self, x):
-        return x, self.xp.zeros(len(x), device=get_device(x))
+        return copy_array(x, xp=self.xp), self.xp.zeros(
+            len(x), device=get_device(x)
+        )
 
     def inverse(self, y):
-        return y, self.xp.zeros(len(y), device=get_device(y))
+        return copy_array(y, xp=self.xp), self.xp.zeros(
+            len(y), device=get_device(y)
+        )
 
 
 class CompositeTransform(BaseTransform):
@@ -329,6 +332,8 @@ class ProbitTransform(BaseTransform):
         return self.forward(x)[0]
 
     def forward(self, x):
+        from scipy.special import erfinv
+
         y = (x - self.lower) / (self.upper - self.lower)
         y = self.xp.clip(y, self.eps, 1.0 - self.eps)
         y = erfinv(2 * y - 1) * math.sqrt(2)
@@ -339,6 +344,8 @@ class ProbitTransform(BaseTransform):
         return y, log_abs_det_jacobian
 
     def inverse(self, y):
+        from scipy.special import erf
+
         log_abs_det_jacobian = (
             -(0.5 * (math.log(2 * math.pi) + y**2)).sum(-1)
             - self._scale_log_abs_det_jacobian