gradbridge 0.1.0__tar.gz

gradbridge-0.1.0/LICENSE

@@ -0,0 +1,15 @@
+                                 Apache License
+                           Version 2.0, January 2004
+                        http://www.apache.org/licenses/
+
+   Licensed under the Apache License, Version 2.0 (the "License");
+   you may not use this file except in compliance with the License.
+   You may obtain a copy of the License at
+
+       http://www.apache.org/licenses/LICENSE-2.0
+
+   Unless required by applicable law or agreed to in writing, software
+   distributed under the License is distributed on an "AS IS" BASIS,
+   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+   See the License for the specific language governing permissions and
+   limitations under the License.

gradbridge-0.1.0/PKG-INFO

@@ -0,0 +1,116 @@
+Metadata-Version: 2.4
+Name: gradbridge
+Version: 0.1.0
+Summary: Gradients that cross framework boundaries: PyTorch, JAX, TensorFlow, NVIDIA Warp and NumPy.
+Author: Anthony Angeles
+License: Apache-2.0
+Project-URL: Repository, https://github.com/aneangel/gradbridge
+Project-URL: Issues, https://github.com/aneangel/gradbridge/issues
+Project-URL: Benchmarks, https://github.com/aneangel/gradbridge/blob/main/benchmarks/RESULTS.md
+Keywords: autodiff,gradients,pytorch,jax,tensorflow,warp,interop,dlpack
+Classifier: Development Status :: 3 - Alpha
+Classifier: License :: OSI Approved :: Apache Software License
+Classifier: Operating System :: POSIX :: Linux
+Classifier: Programming Language :: Python :: 3
+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: Topic :: Scientific/Engineering :: Artificial Intelligence
+Classifier: Intended Audience :: Science/Research
+Requires-Python: >=3.10
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: numpy>=1.24
+Provides-Extra: torch
+Requires-Dist: torch>=2.4; extra == "torch"
+Provides-Extra: jax
+Requires-Dist: jax>=0.4.31; extra == "jax"
+Provides-Extra: tf
+Requires-Dist: tensorflow>=2.15; extra == "tf"
+Provides-Extra: warp
+Requires-Dist: warp-lang>=1.5; extra == "warp"
+Provides-Extra: all
+Requires-Dist: torch>=2.4; extra == "all"
+Requires-Dist: jax>=0.4.31; extra == "all"
+Requires-Dist: tensorflow>=2.15; extra == "all"
+Requires-Dist: warp-lang>=1.5; extra == "all"
+Provides-Extra: test
+Requires-Dist: pytest>=7; extra == "test"
+Dynamic: license-file
+
+# gradbridge
+
+[![ci](https://github.com/aneangel/gradbridge/actions/workflows/ci.yml/badge.svg)](https://github.com/aneangel/gradbridge/actions/workflows/ci.yml)
+
+**Gradients that cross framework boundaries.** Wrap a function written in PyTorch, JAX,
+TensorFlow, NVIDIA Warp or NumPy, then call it — and differentiate through it — from any of
+the others. On the GPU, data stays on the device.
+
+## Getting started
+
+```bash
+pip install "gradbridge[torch,jax]"   # pick the frameworks you use: torch, jax, tf, warp
+pip install "gradbridge[all]"         # or all four
+```
+
+gradbridge itself only requires NumPy and Python 3.10+; each framework adapter is imported
+when you first use it.
+
+## Quick example
+
+Call a JAX function from PyTorch, and backpropagate through it:
+
+```python
+import gradbridge as gb
+import jax.numpy as jnp
+import torch
+
+physics = gb.to_torch(gb.from_jax(lambda x: jnp.sin(x) ** 2))   # JAX code...
+
+x = torch.randn(3, requires_grad=True)
+physics(x).sum().backward()                                     # ...PyTorch gradients
+print(x.grad)                                                   # 2 sin(x) cos(x)
+```
+
+Every combination works the same way: a **source** wraps a function
+(`from_torch`, `from_jax`, `from_tf`, `from_warp`, `from_numpy`) and a **consumer** exposes it
+to another framework (`to_torch`, `to_jax`, `to_tf`, `to_warp`, `to_numpy`). Gradients flow
+through `jit`, `torch.compile`, `vmap` and higher-order derivatives where the frameworks allow.
+
+More complete examples:
+
+- [Warp simulation optimized with JAX](https://github.com/aneangel/gradbridge/blob/main/examples/warp_sim_in_jax.py)
+- [PyTorch policy trained through JAX dynamics and a TensorFlow reward](https://github.com/aneangel/gradbridge/blob/main/examples/mixed_pipeline.py)
+- [Training through a Newton physics simulation from PyTorch](https://github.com/aneangel/gradbridge/blob/main/examples/newton_ball_torch.py)
+  — reproduces Newton's native training loop bit for bit
+
+## Known limitations
+
+- **Bridge coarse functions.** Each call costs a few tenths of a millisecond and nothing fuses
+  across the boundary: wrap a simulator rollout, a model or a loss, not single operations.
+- **NumPy functions run on the host**, so their inputs and outputs are copied off the GPU.
+- **Calling other frameworks from TensorFlow needs a C++ compiler** for zero-copy; without one,
+  gradbridge warns once and falls back to `tf.py_function` (host copies).
+- **Second-order derivatives** need a source that supports them: PyTorch, JAX and TensorFlow do;
+  Warp is first-order only; NumPy uses finite differences beyond your VJP.
+- **Forward mode** is available for the JAX and NumPy consumers only.
+- **Warp buffers you write yourself** between forward and backward are invisible to
+  gradbridge; pass `cache_residuals=False` to `to_torch` in that case.
+- **Import PyTorch (or `gradbridge.torch`) before TensorFlow** if you use `torch.compile`:
+  the two bundle conflicting LLVM builds.
+- **Tested on one GPU** (NVIDIA RTX 5090) and on CPU; multi-GPU placement is not validated yet.
+
+Details and workarounds: [docs/limitations.md](https://github.com/aneangel/gradbridge/blob/main/docs/limitations.md).
+
+## Learn more
+
+- [User guide](https://github.com/aneangel/gradbridge/blob/main/docs/guide.md): what works with
+  each framework, options, NumPy functions, gradient checking, and the `__gradbridge__` protocol
+- [Design](https://github.com/aneangel/gradbridge/blob/main/docs/design.md): how it works, how data
+  moves between frameworks on the GPU, and performance
+- [Benchmarks](https://github.com/aneangel/gradbridge/blob/main/benchmarks/RESULTS.md)
+
+## License
+
+Apache-2.0

gradbridge-0.1.0/README.md

@@ -0,0 +1,75 @@
+# gradbridge
+
+[![ci](https://github.com/aneangel/gradbridge/actions/workflows/ci.yml/badge.svg)](https://github.com/aneangel/gradbridge/actions/workflows/ci.yml)
+
+**Gradients that cross framework boundaries.** Wrap a function written in PyTorch, JAX,
+TensorFlow, NVIDIA Warp or NumPy, then call it — and differentiate through it — from any of
+the others. On the GPU, data stays on the device.
+
+## Getting started
+
+```bash
+pip install "gradbridge[torch,jax]"   # pick the frameworks you use: torch, jax, tf, warp
+pip install "gradbridge[all]"         # or all four
+```
+
+gradbridge itself only requires NumPy and Python 3.10+; each framework adapter is imported
+when you first use it.
+
+## Quick example
+
+Call a JAX function from PyTorch, and backpropagate through it:
+
+```python
+import gradbridge as gb
+import jax.numpy as jnp
+import torch
+
+physics = gb.to_torch(gb.from_jax(lambda x: jnp.sin(x) ** 2))   # JAX code...
+
+x = torch.randn(3, requires_grad=True)
+physics(x).sum().backward()                                     # ...PyTorch gradients
+print(x.grad)                                                   # 2 sin(x) cos(x)
+```
+
+Every combination works the same way: a **source** wraps a function
+(`from_torch`, `from_jax`, `from_tf`, `from_warp`, `from_numpy`) and a **consumer** exposes it
+to another framework (`to_torch`, `to_jax`, `to_tf`, `to_warp`, `to_numpy`). Gradients flow
+through `jit`, `torch.compile`, `vmap` and higher-order derivatives where the frameworks allow.
+
+More complete examples:
+
+- [Warp simulation optimized with JAX](https://github.com/aneangel/gradbridge/blob/main/examples/warp_sim_in_jax.py)
+- [PyTorch policy trained through JAX dynamics and a TensorFlow reward](https://github.com/aneangel/gradbridge/blob/main/examples/mixed_pipeline.py)
+- [Training through a Newton physics simulation from PyTorch](https://github.com/aneangel/gradbridge/blob/main/examples/newton_ball_torch.py)
+  — reproduces Newton's native training loop bit for bit
+
+## Known limitations
+
+- **Bridge coarse functions.** Each call costs a few tenths of a millisecond and nothing fuses
+  across the boundary: wrap a simulator rollout, a model or a loss, not single operations.
+- **NumPy functions run on the host**, so their inputs and outputs are copied off the GPU.
+- **Calling other frameworks from TensorFlow needs a C++ compiler** for zero-copy; without one,
+  gradbridge warns once and falls back to `tf.py_function` (host copies).
+- **Second-order derivatives** need a source that supports them: PyTorch, JAX and TensorFlow do;
+  Warp is first-order only; NumPy uses finite differences beyond your VJP.
+- **Forward mode** is available for the JAX and NumPy consumers only.
+- **Warp buffers you write yourself** between forward and backward are invisible to
+  gradbridge; pass `cache_residuals=False` to `to_torch` in that case.
+- **Import PyTorch (or `gradbridge.torch`) before TensorFlow** if you use `torch.compile`:
+  the two bundle conflicting LLVM builds.
+- **Tested on one GPU** (NVIDIA RTX 5090) and on CPU; multi-GPU placement is not validated yet.
+
+Details and workarounds: [docs/limitations.md](https://github.com/aneangel/gradbridge/blob/main/docs/limitations.md).
+
+## Learn more
+
+- [User guide](https://github.com/aneangel/gradbridge/blob/main/docs/guide.md): what works with
+  each framework, options, NumPy functions, gradient checking, and the `__gradbridge__` protocol
+- [Design](https://github.com/aneangel/gradbridge/blob/main/docs/design.md): how it works, how data
+  moves between frameworks on the GPU, and performance
+- [Benchmarks](https://github.com/aneangel/gradbridge/blob/main/benchmarks/RESULTS.md)
+
+## License
+
+Apache-2.0

gradbridge-0.1.0/pyproject.toml

@@ -0,0 +1,50 @@
+[build-system]
+requires = ["setuptools>=68", "wheel"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "gradbridge"
+version = "0.1.0"
+description = "Gradients that cross framework boundaries: PyTorch, JAX, TensorFlow, NVIDIA Warp and NumPy."
+readme = "README.md"
+requires-python = ">=3.10"
+license = { text = "Apache-2.0" }
+authors = [{ name = "Anthony Angeles" }]
+dependencies = ["numpy>=1.24"]
+keywords = ["autodiff", "gradients", "pytorch", "jax", "tensorflow", "warp", "interop", "dlpack"]
+classifiers = [
+    "Development Status :: 3 - Alpha",
+    "License :: OSI Approved :: Apache Software License",
+    "Operating System :: POSIX :: Linux",
+    "Programming Language :: Python :: 3",
+    "Programming Language :: Python :: 3.10",
+    "Programming Language :: Python :: 3.11",
+    "Programming Language :: Python :: 3.12",
+    "Programming Language :: Python :: 3.13",
+    "Topic :: Scientific/Engineering :: Artificial Intelligence",
+    "Intended Audience :: Science/Research",
+]
+
+[project.urls]
+Repository = "https://github.com/aneangel/gradbridge"
+Issues = "https://github.com/aneangel/gradbridge/issues"
+Benchmarks = "https://github.com/aneangel/gradbridge/blob/main/benchmarks/RESULTS.md"
+
+[project.optional-dependencies]
+torch = ["torch>=2.4"]
+jax = ["jax>=0.4.31"]
+tf = ["tensorflow>=2.15"]
+warp = ["warp-lang>=1.5"]
+all = ["torch>=2.4", "jax>=0.4.31", "tensorflow>=2.15", "warp-lang>=1.5"]
+test = ["pytest>=7"]
+
+[tool.setuptools.packages.find]
+where = ["src"]
+
+[tool.setuptools.package-data]
+# Compiled on first use against the installed TensorFlow (see gradbridge/_tf_native.py).
+gradbridge = ["_tf_op.cc"]
+
+[tool.pytest.ini_options]
+testpaths = ["tests"]
+filterwarnings = ["ignore::DeprecationWarning", "ignore::FutureWarning"]

gradbridge-0.1.0/setup.cfg

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

gradbridge-0.1.0/src/gradbridge/__init__.py

@@ -0,0 +1,129 @@
+"""gradbridge: gradients that cross framework boundaries.
+
+Wrap a function from one framework, use it -- and differentiate through it --
+from any other::
+
+    import gradbridge as gb
+
+    f = gb.from_jax(jax_fn)          # JAX function -> protocol object
+    torch_f = gb.to_torch(f)         # ...callable from PyTorch
+    torch_f(x).sum().backward()      # gradients flow back into JAX
+
+Sources:   from_torch, from_jax, from_tf, from_numpy, from_warp
+Consumers: to_torch,   to_jax,   to_tf,   to_numpy,   to_warp
+
+Framework adapters are imported lazily, so only the frameworks you use need to
+be installed.
+"""
+
+from ._check import CompareReport, GradcheckReport, compare, gradcheck
+from ._convert import HostCopyWarning
+from ._core import DiffFunction, MaxOrderError, Spec, as_diff_function
+
+__version__ = "0.1.0"
+
+__all__ = [
+    "DiffFunction",
+    "Spec",
+    "MaxOrderError",
+    "HostCopyWarning",
+    "as_diff_function",
+    "from_torch",
+    "from_jax",
+    "from_tf",
+    "from_numpy",
+    "from_warp",
+    "to_torch",
+    "to_jax",
+    "to_tf",
+    "to_numpy",
+    "to_warp",
+    "convert",
+    "gradcheck",
+    "compare",
+    "GradcheckReport",
+    "CompareReport",
+]
+
+
+def from_torch(fn, **kwargs):
+    """Wrap a PyTorch function. See :func:`gradbridge.torch.from_torch`."""
+    from .torch import from_torch as _f
+
+    return _f(fn, **kwargs)
+
+
+def from_jax(fn, **kwargs):
+    """Wrap a JAX function. See :func:`gradbridge.jax.from_jax`."""
+    from .jax import from_jax as _f
+
+    return _f(fn, **kwargs)
+
+
+def from_tf(fn, **kwargs):
+    """Wrap a TensorFlow function. See :func:`gradbridge.tf.from_tf`."""
+    from .tf import from_tf as _f
+
+    return _f(fn, **kwargs)
+
+
+def from_numpy(fn, **kwargs):
+    """Wrap a NumPy function. See :func:`gradbridge.numpy.from_numpy`."""
+    from .numpy import from_numpy as _f
+
+    return _f(fn, **kwargs)
+
+
+def from_warp(fn, **kwargs):
+    """Wrap a function that launches Warp kernels. See :func:`gradbridge.warp.from_warp`."""
+    from .warp import from_warp as _f
+
+    return _f(fn, **kwargs)
+
+
+def to_torch(fn, **kwargs):
+    """Expose a gradbridge function to PyTorch. See :func:`gradbridge.torch.to_torch`."""
+    from .torch import to_torch as _f
+
+    return _f(fn, **kwargs)
+
+
+def to_jax(fn, **kwargs):
+    """Expose a gradbridge function to JAX. See :func:`gradbridge.jax.to_jax`."""
+    from .jax import to_jax as _f
+
+    return _f(fn, **kwargs)
+
+
+def to_tf(fn, **kwargs):
+    """Expose a gradbridge function to TensorFlow. See :func:`gradbridge.tf.to_tf`."""
+    from .tf import to_tf as _f
+
+    return _f(fn, **kwargs)
+
+
+def to_numpy(fn, **kwargs):
+    """Expose a gradbridge function to NumPy. See :func:`gradbridge.numpy.to_numpy`."""
+    from .numpy import to_numpy as _f
+
+    return _f(fn, **kwargs)
+
+
+def to_warp(fn, **kwargs):
+    """Expose a gradbridge function to Warp. See :func:`gradbridge.warp.to_warp`."""
+    from .warp import to_warp as _f
+
+    return _f(fn, **kwargs)
+
+
+_SOURCES = {"torch": from_torch, "jax": from_jax, "tf": from_tf, "tensorflow": from_tf, "numpy": from_numpy, "warp": from_warp}
+_CONSUMERS = {"torch": to_torch, "jax": to_jax, "tf": to_tf, "tensorflow": to_tf, "numpy": to_numpy, "warp": to_warp}
+
+
+def convert(fn, *, src: str, dst: str, src_options: dict = None, dst_options: dict = None):
+    """One-step bridge: ``convert(jax_fn, src="jax", dst="torch")``."""
+    try:
+        source, consumer = _SOURCES[src], _CONSUMERS[dst]
+    except KeyError as e:
+        raise ValueError(f"unknown framework {e.args[0]!r}; choose from {sorted(_SOURCES)}") from None
+    return consumer(source(fn, **(src_options or {})), **(dst_options or {}))

gradbridge-0.1.0/src/gradbridge/_check.py

@@ -0,0 +1,148 @@
+"""Gradient checking that works the same way for every framework.
+
+``gradcheck`` compares a function's analytic VJP (any order) against central
+finite differences. ``compare`` checks that two implementations -- say the same
+model in JAX and in PyTorch -- agree on values and gradients.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import List, Optional
+
+import numpy as np
+
+from . import _convert, _fd
+from ._core import as_diff_function
+
+
+@dataclass
+class InputCheck:
+    index: int
+    shape: tuple
+    max_abs_err: float
+    max_rel_err: float
+    ok: bool
+
+
+@dataclass
+class GradcheckReport:
+    name: str
+    order: int
+    inputs: List[InputCheck] = field(default_factory=list)
+    note: str = ""
+
+    @property
+    def passed(self) -> bool:
+        return all(c.ok for c in self.inputs)
+
+    def __bool__(self) -> bool:
+        return self.passed
+
+    def __str__(self) -> str:
+        status = "PASSED" if self.passed else "FAILED"
+        lines = [f"gradcheck '{self.name}' (order {self.order}): {status}"]
+        if self.note:
+            lines.append(f"  {self.note}")
+        lines.append(f"  {'input':>5}  {'shape':<14} {'max abs err':>12} {'max rel err':>12}  ok")
+        for c in self.inputs:
+            lines.append(
+                f"  {c.index:>5}  {str(c.shape):<14} {c.max_abs_err:>12.3e} {c.max_rel_err:>12.3e}  {'yes' if c.ok else 'NO'}"
+            )
+        return "\n".join(lines)
+
+
+def _np(x):
+    return np.asarray(_convert.to_numpy(x))
+
+
+def _random_like(rng, arrays):
+    return [
+        rng.standard_normal(np.shape(a)).astype(a.dtype) if np.issubdtype(a.dtype, np.floating) else np.zeros_like(a)
+        for a in arrays
+    ]
+
+
+def gradcheck(
+    fn,
+    *xs,
+    order: int = 1,
+    eps: Optional[float] = None,
+    rtol: float = 1e-4,
+    atol: float = 1e-6,
+    seed: int = 0,
+    raise_on_fail: bool = False,
+) -> GradcheckReport:
+    """Check ``fn``'s order-``order`` reverse-mode derivative against finite differences.
+
+    Use float64 inputs: float32 finite differences are too noisy for tight
+    tolerances. Order 2 checks the VJP-of-VJP (Hessian-vector products), the
+    same way ``torch.autograd.gradgradcheck`` does.
+    """
+    dfn = as_diff_function(fn)
+    rng = np.random.default_rng(seed)
+    args = [_np(x) for x in xs]
+    target = dfn
+    for _ in range(order - 1):
+        outs = [_np(o) for o in target.eval(*args)]
+        args = args + _random_like(rng, outs)
+        target = target.vjp_function
+
+    outs = [_np(o) for o in target.eval(*args)]
+    cts = _random_like(rng, outs)
+    analytic = [_np(g) for g in target.vjp_function.eval(*args, *cts)]
+    numeric = _fd.fd_vjp(target, *args, cts=cts, eps=eps)
+
+    report = GradcheckReport(name=dfn.name, order=order)
+    if any(a.dtype == np.float32 for a in args):
+        report.note = "float32 inputs: finite differences are noisy, prefer float64"
+    for i, (a, n, x) in enumerate(zip(analytic, numeric, args)):
+        a = a.astype(np.float64).reshape(np.shape(x))
+        n = np.asarray(n, dtype=np.float64)
+        abs_err = float(np.max(np.abs(a - n))) if a.size else 0.0
+        denom = np.maximum(np.abs(n), atol)
+        rel_err = float(np.max(np.abs(a - n) / denom)) if a.size else 0.0
+        ok = bool(np.allclose(a, n, rtol=rtol, atol=atol))
+        report.inputs.append(InputCheck(i, tuple(np.shape(x)), abs_err, rel_err, ok))
+    if raise_on_fail and not report.passed:
+        raise AssertionError(str(report))
+    return report
+
+
+@dataclass
+class CompareReport:
+    name_a: str
+    name_b: str
+    max_value_err: float
+    max_grad_err: float
+    passed: bool
+
+    def __bool__(self) -> bool:
+        return self.passed
+
+    def __str__(self) -> str:
+        status = "MATCH" if self.passed else "MISMATCH"
+        return (
+            f"compare '{self.name_a}' vs '{self.name_b}': {status}  "
+            f"(max value err {self.max_value_err:.3e}, max grad err {self.max_grad_err:.3e})"
+        )
+
+
+def compare(fn_a, fn_b, *xs, rtol: float = 1e-5, atol: float = 1e-7, seed: int = 0) -> CompareReport:
+    """Check that two implementations agree on outputs and on VJPs at ``xs``."""
+    a, b = as_diff_function(fn_a), as_diff_function(fn_b)
+    rng = np.random.default_rng(seed)
+    args = [_np(x) for x in xs]
+    oa = [_np(o) for o in a.eval(*args)]
+    ob = [_np(o) for o in b.eval(*args)]
+    if len(oa) != len(ob) or any(x.shape != y.shape for x, y in zip(oa, ob)):
+        return CompareReport(a.name, b.name, float("inf"), float("inf"), False)
+    val_err = max((float(np.max(np.abs(x - y))) if x.size else 0.0) for x, y in zip(oa, ob)) if oa else 0.0
+    cts = _random_like(rng, oa)
+    ga = [_np(g) for g in a.vjp_function.eval(*args, *cts)]
+    gb = [_np(g) for g in b.vjp_function.eval(*args, *cts)]
+    grad_err = max((float(np.max(np.abs(x - y))) if np.size(x) else 0.0) for x, y in zip(ga, gb)) if ga else 0.0
+    ok = all(np.allclose(x, y, rtol=rtol, atol=atol) for x, y in zip(oa, ob)) and all(
+        np.allclose(x, y, rtol=rtol, atol=atol) for x, y in zip(ga, gb)
+    )
+    return CompareReport(a.name, b.name, val_err, grad_err, ok)