Make torch_interface compatible with torch.func transforms via torch.library.custom_op

Open
  • Is something wrong?
#236 0 comments 0 reactions 0 assignees View on GitHub

Maintainers usually reply within 1 day

Nobody has claimed this yet.

Assessment

Difficulty
5/5
Estimated time
Over a week
Newbie friendliness
35/100
Issue type
Feature
Clarity
Clearly specified
Activity status
Active
Tech stack
python, pytorch, tensorflow
Domain
backend, machine-learning

Research direction

Start with the TODO at tensorcircuit/interfaces/torch.py:14 and inspect the existing torch_interface and its tests in tests/test_interfaces.py and tests/test_torchnn.py. First confirm which torch versions provide the required custom_op, autograd, and vmap APIs; the minimum version is still undecided. Done means torch.func grad, vmap(grad), jacrev, and hessian match native JAX for JAX and TF sources, while classic backward and complex VJP tests still pass.

Written by the indexing model from the issue text.

Description

Problem

tc.interfaces.torch_interface and TorchLayer, which is built on it, wrap a function from another backend in an old-style torch.autograd.Function, where forward(ctx, ...) has no setup_context. Classic loss.backward() works with this wrapper. Every torch.func transform (grad, vjp, vmap, jacrev, jacfwd, hessian) fails on it with:

RuntimeError: In order to use an autograd.Function with functorch transforms (vmap, grad, jvp, jacrev, ...), it must override the setup_context staticmethod.

The torch backend's K.grad is already built on torch.func.grad_and_value, so it fails the same way. K.vjp/K.jvp will also fail once they move to torch.func (#224). Affected use cases include per-sample gradients through functional_call + vmap, and Jacobians or Hessians of hybrid torch models that contain a TorchLayer. This is the existing TODO at tensorcircuit/interfaces/torch.py:14.

Why a new-style autograd.Function is not enough

Moving to the setup_context style does not fix this. Under torch.func, forward receives plain tensors, but backward receives the saved tensors and cotangents as functorch-wrapped tensors (GradTrackingTensor). torch.func does this because it expects backward to be written in torch ops so it can be differentiated again. torch_interface's backward passes those tensors to JAX/TF through dlpack/numpy, and that fails:

RuntimeError: Cannot access data pointer of Tensor that doesn't have storage

A prototype with setup_context plus jvp and vmap staticmethods (JAX source backend, torch 2.2) fails this way under torch.func.grad, vmap(grad), jacrev, jacfwd and hessian. Classic backward still works.

Proposed solution

Use torch.library.custom_op, PyTorch's mechanism for wrapping non-PyTorch code. The implementation receives plain tensors under every transform. A possible design:

  • Forward op: run fun on the source backend.
  • register_autograd: the backward calls a second custom op that runs backend.vjp on the source backend. It keeps the complex-gradient convention handling that vjp_fun does today.
  • register_vmap: map over the source function with backend.vmap(fun) instead of vmapping outside the interface.
  • Higher-order derivatives: also register autograd on the vjp op (vjp-of-vjp on the source backend). This makes hessian work, and so does the reverse-over-reverse K.jvp of the torch backend after #224.
  • Forward mode (torch.func.jvp/jacfwd): custom_op does not seem to support registering a forward rule. This needs checking. The torch backend's K.jvp/K.jacfwd will not need it, because they use reverse-over-reverse.
  • Keep support for pytree inputs/outputs, enable_dlpack, devices, jit, and torch_interface_kws.

Blocker: minimum torch version

torch.library.custom_op and register_autograd need torch >= 2.4. register_vmap may need a newer release; this needs confirming. The repository does not declare a minimum torch version yet (requirements-docker-v2.txt only says torch>2.0). There are two options:

  1. Raise the minimum torch version to one that provides all needed torch.library APIs (simpler, preferred).
  2. Keep the current autograd.Function path for older torch behind a version branch.

This is deferred until the minimum torch version is decided.

Acceptance

  • Results match native JAX for torch.func.grad, vmap(grad), jacrev and hessian through torch_interface, with JAX and TF as source backends.
  • The classic loss.backward() path of TorchLayer and existing tests/test_interfaces.py / tests/test_torchnn.py behavior stay unchanged, including the complex conventions in test_torch_interface_complex_vjp.
Dominant language
Python
Stars
92
Forks
25
Avg merge
21h 12m
Merged PRs (30d)
57

Getting set up

Open in Codespaces

Starts the project's dev container in your browser, under your own GitHub account.

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

More from tensorcircuit/tensorcircuit-ng

All issues in tensorcircuit/tensorcircuit-ng

Similar issues

More Python issues

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.