hoptorch 0.1.0__tar.gz

hoptorch-0.1.0/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2026 Vincent Moens
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

hoptorch-0.1.0/PKG-INFO

@@ -0,0 +1,109 @@
+Metadata-Version: 2.4
+Name: hoptorch
+Version: 0.1.0
+Summary: Compatibility helpers for PyTorch higher-order operators.
+Author: Vincent Moens
+License-Expression: MIT
+Project-URL: Homepage, https://github.com/vmoens/hoptorch
+Project-URL: Repository, https://github.com/vmoens/hoptorch
+Project-URL: Issues, https://github.com/vmoens/hoptorch/issues
+Keywords: pytorch,torch,scan,higher-order-operators,autograd
+Classifier: Development Status :: 3 - Alpha
+Classifier: Intended Audience :: Developers
+Classifier: Intended Audience :: Science/Research
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3 :: Only
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Programming Language :: Python :: 3.13
+Classifier: Programming Language :: Python :: 3.14
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Classifier: Topic :: Software Development :: Libraries :: Python Modules
+Requires-Python: >=3.10
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: torch>=2.7
+Requires-Dist: pyvers>=0.2.2
+Dynamic: license-file
+
+# hoptorch
+
+`hoptorch` is a small compatibility package for PyTorch higher-order operators.
+Its first helper is a safe wrapper around `torch._higher_order_ops.scan` that
+checks whether scan backward is healthy on the requested device and lazily
+installs version-specific compatibility patches for affected PyTorch internals.
+For PyTorch 2.7, where scan backward is not implemented, `hoptorch` backports
+a small eager scan implementation with ordinary autograd support.
+
+## Install
+
+```bash
+pip install hoptorch
+```
+
+Runtime dependency:
+
+```text
+torch>=2.7
+pyvers>=0.2.2
+```
+
+## Usage
+
+```python
+import torch
+from hoptorch import scan
+from hoptorch.scan import ensure_scan_backward, scan_unavailable_reason
+
+if ensure_scan_backward("cpu"):
+    xs = torch.arange(4.0)
+
+    def step(carry, x):
+        next_carry = carry + x
+        return next_carry, next_carry.clone()
+
+    carry, ys = scan(step, torch.zeros(()), xs)
+else:
+    print(scan_unavailable_reason("cpu"))
+```
+
+For compiled code, warm the health check before entering `torch.compile`:
+
+```python
+from hoptorch.scan import ensure_scan_backward
+
+ensure_scan_backward("cpu")
+compiled_fn = torch.compile(fn)
+```
+
+If a wrapper call happens during Dynamo tracing before the health result is
+cached, `hoptorch` fails closed instead of tracing the probe.
+
+## Public API
+
+```python
+from hoptorch import scan
+from hoptorch.scan import (
+    ensure_scan_backward,
+    has_scan,
+    patch_scan_backward,
+    scan_unavailable_reason,
+)
+```
+
+- `scan(fn, init, xs, *, dim=0, **kwargs)`: calls PyTorch scan only when scan
+  backward is known to be healthy for the inferred device.
+- `has_scan()`: reports whether `torch._higher_order_ops.scan.scan` exists.
+- `ensure_scan_backward(device=None)`: runs or reads the cached health check,
+  patching lazily if needed, and returns `True` only after a passing probe.
+- `scan_unavailable_reason(device=None)`: returns `None` when scan backward is
+  usable, otherwise a stable human-readable reason.
+- `patch_scan_backward()`: attempts to install the compatibility patch and is
+  intended mainly for diagnostics.
+
+## Tests
+
+```bash
+python -m unittest discover -s tests
+```

hoptorch-0.1.0/README.md

@@ -0,0 +1,80 @@
+# hoptorch
+
+`hoptorch` is a small compatibility package for PyTorch higher-order operators.
+Its first helper is a safe wrapper around `torch._higher_order_ops.scan` that
+checks whether scan backward is healthy on the requested device and lazily
+installs version-specific compatibility patches for affected PyTorch internals.
+For PyTorch 2.7, where scan backward is not implemented, `hoptorch` backports
+a small eager scan implementation with ordinary autograd support.
+
+## Install
+
+```bash
+pip install hoptorch
+```
+
+Runtime dependency:
+
+```text
+torch>=2.7
+pyvers>=0.2.2
+```
+
+## Usage
+
+```python
+import torch
+from hoptorch import scan
+from hoptorch.scan import ensure_scan_backward, scan_unavailable_reason
+
+if ensure_scan_backward("cpu"):
+    xs = torch.arange(4.0)
+
+    def step(carry, x):
+        next_carry = carry + x
+        return next_carry, next_carry.clone()
+
+    carry, ys = scan(step, torch.zeros(()), xs)
+else:
+    print(scan_unavailable_reason("cpu"))
+```
+
+For compiled code, warm the health check before entering `torch.compile`:
+
+```python
+from hoptorch.scan import ensure_scan_backward
+
+ensure_scan_backward("cpu")
+compiled_fn = torch.compile(fn)
+```
+
+If a wrapper call happens during Dynamo tracing before the health result is
+cached, `hoptorch` fails closed instead of tracing the probe.
+
+## Public API
+
+```python
+from hoptorch import scan
+from hoptorch.scan import (
+    ensure_scan_backward,
+    has_scan,
+    patch_scan_backward,
+    scan_unavailable_reason,
+)
+```
+
+- `scan(fn, init, xs, *, dim=0, **kwargs)`: calls PyTorch scan only when scan
+  backward is known to be healthy for the inferred device.
+- `has_scan()`: reports whether `torch._higher_order_ops.scan.scan` exists.
+- `ensure_scan_backward(device=None)`: runs or reads the cached health check,
+  patching lazily if needed, and returns `True` only after a passing probe.
+- `scan_unavailable_reason(device=None)`: returns `None` when scan backward is
+  usable, otherwise a stable human-readable reason.
+- `patch_scan_backward()`: attempts to install the compatibility patch and is
+  intended mainly for diagnostics.
+
+## Tests
+
+```bash
+python -m unittest discover -s tests
+```

hoptorch-0.1.0/pyproject.toml

@@ -0,0 +1,39 @@
+[build-system]
+requires = ["setuptools>=68", "wheel"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "hoptorch"
+version = "0.1.0"
+description = "Compatibility helpers for PyTorch higher-order operators."
+readme = "README.md"
+requires-python = ">=3.10"
+license = "MIT"
+authors = [{ name = "Vincent Moens" }]
+keywords = ["pytorch", "torch", "scan", "higher-order-operators", "autograd"]
+classifiers = [
+    "Development Status :: 3 - Alpha",
+    "Intended Audience :: Developers",
+    "Intended Audience :: Science/Research",
+    "Programming Language :: Python :: 3",
+    "Programming Language :: Python :: 3 :: Only",
+    "Programming Language :: Python :: 3.10",
+    "Programming Language :: Python :: 3.11",
+    "Programming Language :: Python :: 3.12",
+    "Programming Language :: Python :: 3.13",
+    "Programming Language :: Python :: 3.14",
+    "Topic :: Scientific/Engineering :: Artificial Intelligence",
+    "Topic :: Software Development :: Libraries :: Python Modules",
+]
+dependencies = ["torch>=2.7", "pyvers>=0.2.2"]
+
+[project.urls]
+Homepage = "https://github.com/vmoens/hoptorch"
+Repository = "https://github.com/vmoens/hoptorch"
+Issues = "https://github.com/vmoens/hoptorch/issues"
+
+[tool.setuptools]
+package-dir = {"" = "src"}
+
+[tool.setuptools.packages.find]
+where = ["src"]

hoptorch-0.1.0/setup.cfg

@@ -0,0 +1,4 @@
+[egg_info]
+tag_build = 
+tag_date = 0
+

hoptorch-0.1.0/src/hoptorch/__init__.py

@@ -0,0 +1,8 @@
+"""Compatibility helpers for PyTorch higher-order operators."""
+
+from __future__ import annotations
+
+from .scan import scan
+
+__all__ = ["scan"]
+__version__ = "0.1.0"

hoptorch-0.1.0/src/hoptorch/_scan_backport_27.py

@@ -0,0 +1,144 @@
+"""PyTorch 2.7 scan backport.
+
+PyTorch 2.7 exposes ``torch._higher_order_ops.scan.scan`` but registers its
+Autograd dispatch key as ``autograd_not_implemented``. For that version, patch
+the public scan function to a small eager Python implementation whose backward
+is handled by ordinary PyTorch autograd through the unrolled loop.
+"""
+
+from __future__ import annotations
+
+from typing import Any, Callable
+
+
+def _canonicalize_dim(ndim: int, dim: int) -> int:
+    if not isinstance(dim, int):
+        raise RuntimeError(f"Dim must be an int, but got {type(dim)}")
+    if ndim == 0:
+        raise RuntimeError("Cannot scan over a scalar xs tensor")
+    if dim < 0:
+        dim += ndim
+    if dim < 0 or dim >= ndim:
+        raise IndexError(f"Dimension out of range (expected 0 <= dim < {ndim}, got {dim})")
+    return dim
+
+
+def _validate_tensor_leaves(torch: Any, leaves: list[Any], name: str) -> None:
+    if not leaves:
+        raise RuntimeError(f"scan() operator requires {name} leaves.")
+    for leaf in leaves:
+        if not isinstance(leaf, torch.Tensor):
+            raise RuntimeError(f"All {name} leaves must be tensors, but got {leaf!r}")
+
+
+def scan_27_backport(
+    combine_fn: Callable[[Any, Any], tuple[Any, Any]],
+    init: Any,
+    xs: Any,
+    *,
+    dim: int = 0,
+    reverse: bool = False,
+) -> tuple[Any, Any]:
+    """Eager scan implementation for PyTorch 2.7 with normal autograd support."""
+
+    import torch
+    import torch.utils._pytree as pytree
+
+    if not callable(combine_fn):
+        raise RuntimeError(f"Combine_fn must be callable, but got {combine_fn!r}")
+    if not isinstance(reverse, bool):
+        raise RuntimeError(f"Reverse must be a bool, but got {type(reverse)}")
+
+    leaves_init, spec_init = pytree.tree_flatten(init)
+    leaves_xs_orig, spec_xs = pytree.tree_flatten(xs)
+    if not leaves_xs_orig:
+        return init, []
+
+    _validate_tensor_leaves(torch, leaves_init, "init")
+    _validate_tensor_leaves(torch, leaves_xs_orig, "xs")
+
+    ndim = leaves_xs_orig[0].ndim
+    dim = _canonicalize_dim(ndim, dim)
+    scan_length = leaves_xs_orig[0].shape[dim]
+    for leaf in leaves_xs_orig:
+        if leaf.ndim <= dim:
+            raise RuntimeError(
+                f"All xs leaves must have dimension {dim}, but got shape {tuple(leaf.shape)}"
+            )
+        if leaf.shape[dim] != scan_length:
+            raise RuntimeError("All xs leaves must have the same scan dimension size")
+
+    leaves_xs = [torch.movedim(leaf, dim, 0) for leaf in leaves_xs_orig]
+    if reverse:
+        leaves_xs = [torch.flip(leaf, [0]) for leaf in leaves_xs]
+
+    carry = init
+    out_spec = None
+    flat_outputs_by_step: list[list[Any]] = []
+    for index in range(scan_length):
+        x_slice = pytree.tree_unflatten(
+            [leaf.select(0, index) for leaf in leaves_xs], spec_xs
+        )
+        carry, output = combine_fn(carry, x_slice)
+        flat_output, current_out_spec = pytree.tree_flatten(output)
+        if out_spec is None:
+            out_spec = current_out_spec
+        elif current_out_spec != out_spec:
+            raise RuntimeError("scan output pytree structure changed across iterations")
+        for leaf in flat_output:
+            if not isinstance(leaf, torch.Tensor):
+                raise RuntimeError(
+                    "hoptorch's PyTorch 2.7 scan backport only supports tensor output leaves"
+                )
+        flat_outputs_by_step.append(flat_output)
+
+    flat_carry, current_carry_spec = pytree.tree_flatten(carry)
+    if current_carry_spec != spec_init:
+        raise RuntimeError("scan carry pytree structure must match init")
+    for init_leaf, carry_leaf in zip(leaves_init, flat_carry):
+        if not isinstance(carry_leaf, torch.Tensor):
+            raise RuntimeError("All carry leaves must be tensors")
+        if init_leaf.shape != carry_leaf.shape or init_leaf.dtype != carry_leaf.dtype:
+            raise RuntimeError("scan carry tensor metadata must match init")
+
+    if out_spec is None:
+        return carry, []
+
+    flat_stacked_outputs = []
+    for output_index in range(len(flat_outputs_by_step[0])):
+        leaves = [step[output_index] for step in flat_outputs_by_step]
+        stacked = torch.stack(leaves, dim=0)
+        if reverse:
+            stacked = torch.flip(stacked, [0])
+        flat_stacked_outputs.append(stacked)
+
+    return carry, pytree.tree_unflatten(flat_stacked_outputs, out_spec)
+
+
+def install_scan_27_backport(scan_module: Any) -> bool:
+    """Patch ``torch._higher_order_ops.scan.scan`` for PyTorch 2.7."""
+
+    if getattr(scan_module, "_hoptorch_scan_27_backport", False):
+        return True
+    original_scan = getattr(scan_module, "scan", None)
+    if original_scan is None:
+        return False
+    scan_27_backport._hoptorch_scan_backward_original = original_scan  # type: ignore[attr-defined]
+    scan_module.scan = scan_27_backport
+    scan_module._hoptorch_scan_27_backport = True
+    return True
+
+
+def rollback_scan_27_backport(scan_module: Any) -> bool:
+    """Restore the original PyTorch 2.7 scan function if this patch was applied."""
+
+    if not getattr(scan_module, "_hoptorch_scan_27_backport", False):
+        return False
+    original_scan = getattr(
+        getattr(scan_module, "scan", None), "_hoptorch_scan_backward_original", None
+    )
+    if original_scan is None:
+        return False
+    scan_module.scan = original_scan
+    scan_module._hoptorch_scan_27_backport = False
+    return True

hoptorch-0.1.0/src/hoptorch/_scan_patch.py

@@ -0,0 +1,251 @@
+"""Private PyTorch monkey patches for scan backward compatibility.
+
+All direct interaction with PyTorch scan autograd internals is isolated here so
+callers do not need to import or reason about those private symbols.
+"""
+
+from __future__ import annotations
+
+from typing import Any
+
+from pyvers import implement_for
+
+
+def _differentiable_sample_requires_grad(value: Any) -> Any:
+    import torch
+
+    if not isinstance(value, torch.Tensor):
+        return value
+    sample = value.detach().clone()
+    if value.dtype.is_floating_point or value.dtype.is_complex:
+        sample.requires_grad_(True)
+    return sample
+
+
+def _map_differentiable_samples(value: Any) -> Any:
+    import torch
+
+    if isinstance(value, torch.Tensor):
+        return _differentiable_sample_requires_grad(value)
+    if isinstance(value, tuple):
+        return tuple(_map_differentiable_samples(item) for item in value)
+    if isinstance(value, list):
+        return [_map_differentiable_samples(item) for item in value]
+    if isinstance(value, dict):
+        return {key: _map_differentiable_samples(item) for key, item in value.items()}
+    return value
+
+
+def _register_scan_autograd_impl(scan_module: Any, scan_autograd: Any) -> bool:
+    import torch
+
+    scan_op = getattr(scan_module, "scan_op", None)
+    if scan_op is None or not hasattr(scan_op, "py_autograd_impl"):
+        return False
+
+    py_kernels = getattr(scan_op, "py_kernels", None)
+    if py_kernels is None:
+        return False
+
+    autograd_key = torch._C.DispatchKey.Autograd
+    previous = py_kernels.pop(autograd_key, None)
+    if hasattr(scan_op, "_dispatch_cache"):
+        scan_op._dispatch_cache.clear()
+    try:
+        scan_op.py_autograd_impl(scan_autograd)
+    except Exception:
+        py_kernels.pop(autograd_key, None)
+        if previous is not None:
+            py_kernels[autograd_key] = previous
+        if hasattr(scan_op, "_dispatch_cache"):
+            scan_op._dispatch_cache.clear()
+        return False
+
+    scan_module.scan_autograd = scan_autograd
+    return True
+
+
+def _patch_backward_input_output_aliasing(scan_module: Any) -> bool:
+    import torch
+    import torch.fx
+
+    try:
+        from torch.multiprocessing.reductions import StorageWeakRef
+    except Exception:
+        StorageWeakRef = None
+
+    scan_autograd_impl = getattr(scan_module, "ScanAutogradImpl", None)
+    if scan_autograd_impl is None:
+        return False
+    if getattr(scan_autograd_impl, "_hoptorch_scan_backward_alias_patch", False):
+        return True
+
+    def _break_bw_input_output_aliasing(self) -> None:
+        if StorageWeakRef is None or not hasattr(self, "_insert_clone"):
+            return
+
+        bw_gm = self.hop_partitioned_graph.bw_gm
+        bw_output_node = next(iter(bw_gm.graph.find_nodes(op="output")))
+        if len(bw_output_node.args) != 1:
+            raise AssertionError(
+                "expected bw_gm output to have 1 arg, got "
+                f"{len(bw_output_node.args)}"
+            )
+
+        bw_outputs = bw_output_node.args[0]
+        if not isinstance(bw_outputs, (tuple, list)):
+            bw_outputs = (bw_outputs,)
+
+        placeholder_storages = set()
+        for placeholder in bw_gm.graph.find_nodes(op="placeholder"):
+            val = placeholder.meta.get("val", None) if hasattr(placeholder, "meta") else None
+            if isinstance(val, torch.Tensor):
+                placeholder_storages.add(StorageWeakRef(val._typed_storage()))
+
+        def aliases_placeholder(node: torch.fx.Node) -> bool:
+            if node.op == "placeholder":
+                return True
+            val = node.meta.get("val", None) if hasattr(node, "meta") else None
+            if isinstance(val, torch.Tensor):
+                return StorageWeakRef(val._typed_storage()) in placeholder_storages
+            return False
+
+        new_bw_outputs = []
+        rewrote = False
+        for output in bw_outputs:
+            if isinstance(output, torch.fx.Node) and aliases_placeholder(output):
+                new_bw_outputs.append(self._insert_clone(output, bw_output_node))
+                rewrote = True
+            else:
+                new_bw_outputs.append(output)
+
+        if rewrote:
+            bw_output_node.args = (tuple(new_bw_outputs),)
+            bw_gm.graph.lint()
+            bw_gm.recompile()
+
+    original_init = scan_autograd_impl.__init__
+
+    def __init__(self, *args: Any, **kwargs: Any) -> None:
+        original_init(self, *args, **kwargs)
+        self._break_bw_input_output_aliasing()
+
+    scan_autograd_impl._break_bw_input_output_aliasing = _break_bw_input_output_aliasing
+    scan_autograd_impl.__init__ = __init__
+    scan_autograd_impl._hoptorch_scan_backward_alias_patch = True
+    return True
+
+
+def _patch_current_scan_autograd(scan_module: Any) -> bool:
+    required = (
+        "HopGraphMinCutPartitioner",
+        "ScanAutogradImpl",
+        "ScanAutogradOp",
+        "disable_proxy_modes_tracing",
+    )
+    if not all(hasattr(scan_module, name) for name in required):
+        return False
+
+    _patch_backward_input_output_aliasing(scan_module)
+
+    def scan_autograd(combine_fn, init, xs, additional_inputs):
+        with scan_module.disable_proxy_modes_tracing():
+            sample_init = [_differentiable_sample_requires_grad(t) for t in init]
+            sample_args = (*sample_init, *[x[0] for x in xs], *additional_inputs)
+            try:
+                hop_partitioned_graph = (
+                    scan_module.HopGraphMinCutPartitioner.create_partitioned_graph(
+                        combine_fn,
+                        sample_args,
+                        always_recompute_complex_exprs=True,
+                    )
+                )
+            except TypeError:
+                real_args = (*init, *[x[0] for x in xs], *additional_inputs)
+                hop_partitioned_graph = (
+                    scan_module.HopGraphMinCutPartitioner.create_partitioned_graph(
+                        combine_fn,
+                        real_args,
+                        sample_args,
+                        always_recompute_complex_exprs=True,
+                    )
+                )
+
+        return scan_module.ScanAutogradOp.apply(
+            hop_partitioned_graph,
+            len(init),
+            len(xs),
+            len(additional_inputs),
+            *init,
+            *xs,
+            *additional_inputs,
+        )
+
+    return _register_scan_autograd_impl(scan_module, scan_autograd)
+
+
+def _patch_older_scan_autograd(scan_module: Any) -> bool:
+    materialize_as_graph = getattr(scan_module, "materialize_as_graph", None)
+    if materialize_as_graph is None:
+        return False
+    if getattr(materialize_as_graph, "_hoptorch_scan_backward_sample_patch", False):
+        return True
+
+    def patched_materialize_as_graph(fn, args, *other_args, **kwargs):
+        return materialize_as_graph(
+            fn,
+            _map_differentiable_samples(args),
+            *other_args,
+            **kwargs,
+        )
+
+    patched_materialize_as_graph._hoptorch_scan_backward_sample_patch = True
+    patched_materialize_as_graph._hoptorch_scan_backward_original = materialize_as_graph
+    scan_module.materialize_as_graph = patched_materialize_as_graph
+    return True
+
+
+@implement_for("torch")
+def install_scan_backward_patch(scan_module: Any) -> bool:
+    """Attempt to install a scan backward compatibility patch."""
+
+    return False
+
+
+@install_scan_backward_patch.register(from_version="2.7", to_version="2.8")
+def _(scan_module: Any) -> bool:
+    from ._scan_backport_27 import install_scan_27_backport
+
+    return install_scan_27_backport(scan_module)
+
+
+@install_scan_backward_patch.register(from_version="2.8")
+def _(scan_module: Any) -> bool:
+    try:
+        return _patch_current_scan_autograd(scan_module) or _patch_older_scan_autograd(scan_module)
+    except Exception:
+        return False
+
+
+@implement_for("torch")
+def rollback_failed_scan_backward_patch(scan_module: Any) -> bool:
+    """Best-effort rollback for reversible compatibility patches."""
+
+    return False
+
+
+@rollback_failed_scan_backward_patch.register(from_version="2.7", to_version="2.8")
+def _(scan_module: Any) -> bool:
+    from ._scan_backport_27 import rollback_scan_27_backport
+
+    return rollback_scan_27_backport(scan_module)
+
+
+@rollback_failed_scan_backward_patch.register(from_version="2.8")
+def _(scan_module: Any) -> bool:
+    materialize_as_graph = getattr(scan_module, "materialize_as_graph", None)
+    original = getattr(materialize_as_graph, "_hoptorch_scan_backward_original", None)
+    if original is None:
+        return False
+    scan_module.materialize_as_graph = original
+    return True