simplify: collapse transform walkers to a single _apply_to_transform (#1893)

#1896 shipped two near-identical tree-walkers (_apply_to_transform and a
test-only _transform_tensors collector). Keep a single self-contained
_apply_to_transform in production and collect tensors via a test helper that
reuses it, removing the test-only function and the visitor indirection from the
production module. Behavior-preserving; device/dtype/walker tests still green.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

Author: janfb

sbi/utils/sbiutils.py

@@ -811,19 +811,19 @@ def match_theta_and_x_batch_shapes(theta: Tensor, x: Tensor) -> Tuple[Tensor, Te
     return theta_repeated, x_repeated
 
 
-def _walk_transform_tensors(
-    transform: TorchTransform, visit: Callable[[object, str, Tensor], None]
+def _apply_to_transform(
+    transform: TorchTransform, fn: Callable[[Tensor], Tensor]
 ) -> None:
-    """Walk a transform tree, calling ``visit(holder, key, tensor)`` per tensor.
+    """Apply fn to all tensors in a transform tree, in place.
 
-    Traverses ``ComposeTransform.parts``, ``IndependentTransform.base_transform``,
-    etc., so callers can move (``.to()``) or collect the tensors held inside a
-    transform, which are otherwise invisible to ``nn.Module._apply``.
+    So ``.to()``/``.double()`` calls propagate into a transform held as a plain
+    attribute (e.g. a prior transform on an estimator), which is otherwise
+    invisible to ``nn.Module._apply``. Traverses ``ComposeTransform.parts``,
+    ``IndependentTransform.base_transform`` and nested inverses.
 
     Args:
         transform: Root of the transform tree.
-        visit: Callback invoked for each tensor with its holder object and
-            attribute name, e.g. to reassign or accumulate it.
+        fn: Callable applied to each tensor (e.g. ``lambda t: t.to(device)``).
     """
     seen = set()
 
@@ -833,7 +833,7 @@ def _walk(t):
         seen.add(id(t))
         for key, val in list(t.__dict__.items()):
             if isinstance(val, Tensor):
-                visit(t, key, val)
+                object.__setattr__(t, key, fn(val))
             elif isinstance(val, (list, tuple)):
                 for item in val:
                     if isinstance(item, TorchTransform):
@@ -844,41 +844,6 @@ def _walk(t):
     _walk(transform)
 
 
-def _apply_to_transform(
-    transform: TorchTransform, fn: Callable[[Tensor], Tensor]
-) -> None:
-    """Apply fn to all tensors in a transform tree.
-
-    So ``.to()``/``.double()`` calls propagate into a transform held as a plain
-    attribute (e.g. a prior transform on an estimator).
-
-    Args:
-        transform: Root of the transform tree.
-        fn: Callable applied to each tensor (e.g. ``lambda t: t.to(device)``).
-    """
-
-    def _move(holder, key, tensor):
-        object.__setattr__(holder, key, fn(tensor))
-
-    _walk_transform_tensors(transform, _move)
-
-
-def _transform_tensors(transform: TorchTransform) -> List[Tensor]:
-    """Collect every tensor in a transform tree.
-
-    Args:
-        transform: Root of the transform tree.
-
-    Returns:
-        List of all tensors found in the transform tree.
-    """
-    tensors: List[Tensor] = []
-    _walk_transform_tensors(
-        transform, lambda _holder, _key, tensor: tensors.append(tensor)
-    )
-    return tensors
-
-
 def mcmc_transform(
     prior: Distribution,
     num_prior_samples_for_zscoring: int = 1000,

tests/inference_on_device_test.py

@@ -58,6 +58,25 @@
 )
 
 
+def _collect_transform_tensors(transform):
+    """Collect every tensor in a transform tree, reusing the production walk.
+
+    Test helper for the transform_to_unconstrained device/dtype checks: drives
+    ``_apply_to_transform`` with a recording identity fn so we don't duplicate the
+    tree-walk in the tests.
+    """
+    from sbi.utils.sbiutils import _apply_to_transform
+
+    collected = []
+
+    def _record(tensor):
+        collected.append(tensor)
+        return tensor
+
+    _apply_to_transform(transform, _record)
+    return collected
+
+
 @pytest.mark.slow
 @pytest.mark.gpu
 @pytest.mark.parametrize(
@@ -897,15 +916,14 @@ def test_npe_pfn_on_device(prior_device):
 def test_mdn_device_transform():
     """MDN with transform_to_unconstrained moves transform tensors on .to()."""
     from sbi.neural_nets.net_builders.mdn import build_mdn
-    from sbi.utils.sbiutils import _transform_tensors
 
     device = process_device("gpu")
     prior = BoxUniform(-2 * torch.ones(2), 2 * torch.ones(2))
     bx, by = prior.sample((512,)), torch.randn(512, 3)
     est = build_mdn(bx, by, z_score_x="transform_to_unconstrained", x_dist=prior)
     est.to(device)
 
-    transform_tensors = _transform_tensors(est._prior_transform)
+    transform_tensors = _collect_transform_tensors(est._prior_transform)
     assert transform_tensors, "expected the prior transform to hold tensors"
     for t in transform_tensors:
         assert t.device.type == device.split(":")[0], (
@@ -928,15 +946,14 @@ def test_mdn_transform_follows_dtype():
     Runs on every CI (no GPU needed).
     """
     from sbi.neural_nets.net_builders.mdn import build_mdn
-    from sbi.utils.sbiutils import _transform_tensors
 
     prior = BoxUniform(-2 * torch.ones(2), 2 * torch.ones(2))
     bx, by = prior.sample((256,)), torch.randn(256, 3)
     est = build_mdn(bx, by, z_score_x="transform_to_unconstrained", x_dist=prior)
 
     est.double()
 
-    transform_tensors = _transform_tensors(est._prior_transform)
+    transform_tensors = _collect_transform_tensors(est._prior_transform)
     assert transform_tensors, "expected the prior transform to hold tensors"
     assert all(t.dtype == torch.float64 for t in transform_tensors)
 
@@ -950,15 +967,14 @@ def test_mdn_transform_follows_dtype():
 def test_zuko_device_transform():
     """Zuko flow with transform_to_unconstrained moves transform tensors on .to()."""
     from sbi.neural_nets.net_builders.flow import build_zuko_maf
-    from sbi.utils.sbiutils import _transform_tensors
 
     device = process_device("gpu")
     prior = BoxUniform(-2 * torch.ones(2), 2 * torch.ones(2))
     bx, by = prior.sample((512,)), torch.randn(512, 3)
     est = build_zuko_maf(bx, by, z_score_x="transform_to_unconstrained", x_dist=prior)
     est.to(device)
 
-    transform_tensors = _transform_tensors(est._prior_transform)
+    transform_tensors = _collect_transform_tensors(est._prior_transform)
     assert transform_tensors, "expected the prior transform to hold tensors"
     for t in transform_tensors:
         assert t.device.type == device.split(":")[0], (
@@ -982,14 +998,13 @@ def test_zuko_transform_follows_dtype():
     .double()/.to() cast would silently leave it on the old dtype/device.
     """
     from sbi.neural_nets.net_builders.flow import build_zuko_maf
-    from sbi.utils.sbiutils import _transform_tensors
 
     prior = BoxUniform(-2 * torch.ones(2), 2 * torch.ones(2))
     bx, by = prior.sample((256,)), torch.randn(256, 3)
     est = build_zuko_maf(bx, by, z_score_x="transform_to_unconstrained", x_dist=prior)
 
     est.double()
 
-    transform_tensors = _transform_tensors(est._prior_transform)
+    transform_tensors = _collect_transform_tensors(est._prior_transform)
     assert transform_tensors, "expected the prior transform to hold tensors"
     assert all(t.dtype == torch.float64 for t in transform_tensors)

tests/sbiutils_test.py

@@ -711,37 +711,29 @@ def test_mdn_transform_to_unconstrained():
     assert s.shape[0] == 10 and torch.isfinite(s).all()
 
 
-def test_walk_transform_tensors_reaches_nested_tensors():
-    """The shared transform walker reaches tensors behind lists, nesting and .inv.
+def test_apply_to_transform_reaches_nested_tensors():
+    """`_apply_to_transform` reaches tensors behind lists, nesting and .inv.
 
-    Guards `_walk_transform_tensors` (and its `_apply_to_transform` /
-    `_transform_tensors` wrappers), which move a prior transform onto a device.
     The leaf `loc`/`scale` sit behind `_InverseTransform -> IndependentTransform ->
     ComposeTransform.parts (a list) -> AffineTransform`, so a walk that missed the
-    list branch or the inverse would silently move nothing.
+    list branch or the inverse would silently touch nothing. A dtype cast is the
+    observable effect: both leaves must change.
     """
     from torch.distributions.transforms import (
         AffineTransform,
         ComposeTransform,
         SigmoidTransform,
     )
 
-    from sbi.utils.sbiutils import _apply_to_transform, _transform_tensors
+    from sbi.utils.sbiutils import _apply_to_transform
 
-    loc, scale = torch.zeros(3), torch.ones(3)
-    affine = AffineTransform(loc, scale)
+    affine = AffineTransform(torch.zeros(3), torch.ones(3))
     # matches the shape mcmc_transform produces for a bounded prior, then inverted.
     transform = IndependentTransform(
         ComposeTransform([SigmoidTransform(), affine]), 1
     ).inv
 
-    # collector reaches exactly the two leaf tensors (same objects), not zero.
-    tensors = _transform_tensors(transform)
-    assert len(tensors) == 2
-    assert any(t is loc for t in tensors)
-    assert any(t is scale for t in tensors)
-
-    # apply reaches every tensor via the same walk (dtype cast as an observable op).
+    assert affine.loc.dtype == torch.float32 and affine.scale.dtype == torch.float32
     _apply_to_transform(transform, lambda t: t.double())
     assert affine.loc.dtype == torch.float64
     assert affine.scale.dtype == torch.float64