Skip to content

Pytorch triple backward #200

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
vbharadwaj-bk merged 10 commits into from
Jun 7, 2026
Merged

Pytorch triple backward #200

Show file tree
Hide file tree
Changes from 3 commits
Commits
Show all changes
10 commits
Select commit Hold shift + click to select a range
6eb39ec
Add PyTorch triple backward support
asglover May 21, 2026
c0d5712
better tests, alleged fix for bad deterministic test
asglover Jun 2, 2026
f54d389
ruff
asglover Jun 2, 2026
f588795
Merge remote-tracking branch 'origin/main' into pytorch-triple-backward
asglover Jun 2, 2026
8efe3ba
put device, requries_grad in constructor directly
asglover Jun 5, 2026
e693f14
remove stream test changes
asglover Jun 5, 2026
82ec199
remove DirectOp tests
asglover Jun 5, 2026
0f139a3
format
asglover Jun 5, 2026
0e5c072
no diff
asglover Jun 5, 2026
23a8d41
remove dead code
asglover Jun 5, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
184 changes: 184 additions & 0 deletions openequivariance/openequivariance/_torch/NPDoubleBackwardMixin.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,6 +1,12 @@
import torch


def _none_to_zeros(values, refs):
return tuple(
ref * 0 if value is None else value for value, ref in zip(values, refs)
)


class NumpyDoubleBackwardMixin:
"""
Adds a Numpy double backward method to any TensorProduct
Expand Down Expand Up @@ -43,6 +49,89 @@ def double_backward_cpu(
d.detach().cpu().numpy(),
)

def triple_backward_cpu(
self,
in1,
in2,
out_grad,
weights,
weights_dgrad,
in1_dgrad,
in2_dgrad,
out_tgrad,
weights_tgrad,
in1_tgrad,
in2_tgrad,
):
assert self.torch_op

in1_torch = torch.tensor(in1).to("cuda").requires_grad_(True)
in2_torch = torch.tensor(in2).to("cuda").requires_grad_(True)
weights_torch = torch.tensor(weights).to("cuda").requires_grad_(True)
out_grad_torch = torch.tensor(out_grad).to("cuda").requires_grad_(True)
in1_dgrad_torch = torch.tensor(in1_dgrad).to("cuda").requires_grad_(True)
in2_dgrad_torch = torch.tensor(in2_dgrad).to("cuda").requires_grad_(True)
weights_dgrad_torch = (
torch.tensor(weights_dgrad).to("cuda").requires_grad_(True)
)
out_tgrad_torch = torch.tensor(out_tgrad).to("cuda")
in1_tgrad_torch = torch.tensor(in1_tgrad).to("cuda")
in2_tgrad_torch = torch.tensor(in2_tgrad).to("cuda")
weights_tgrad_torch = torch.tensor(weights_tgrad).to("cuda")

out_torch = self.forward(in1_torch, in2_torch, weights_torch)
in1_grad, in2_grad, weights_grad = torch.autograd.grad(
outputs=out_torch,
inputs=[in1_torch, in2_torch, weights_torch],
grad_outputs=out_grad_torch,
create_graph=True,
retain_graph=True,
)
double_grads = torch.autograd.grad(
outputs=[in1_grad, in2_grad, weights_grad],
inputs=[in1_torch, in2_torch, weights_torch, out_grad_torch],
grad_outputs=[in1_dgrad_torch, in2_dgrad_torch, weights_dgrad_torch],
create_graph=True,
retain_graph=True,
allow_unused=True,
)
double_grads = _none_to_zeros(
double_grads, (in1_torch, in2_torch, weights_torch, out_grad_torch)
)
triple_grads = torch.autograd.grad(
outputs=double_grads,
inputs=[
in1_torch,
in2_torch,
weights_torch,
out_grad_torch,
in1_dgrad_torch,
in2_dgrad_torch,
weights_dgrad_torch,
],
grad_outputs=[
in1_tgrad_torch,
in2_tgrad_torch,
weights_tgrad_torch,
out_tgrad_torch,
],
allow_unused=True,
)
triple_grads = _none_to_zeros(
triple_grads,
(
in1_torch,
in2_torch,
weights_torch,
out_grad_torch,
in1_dgrad_torch,
in2_dgrad_torch,
weights_dgrad_torch,
),
)

return tuple(grad.detach().cpu().numpy() for grad in triple_grads)


class NumpyDoubleBackwardMixinConv:
"""
Expand Down Expand Up @@ -95,3 +184,98 @@ def double_backward_cpu(
c.detach().cpu().numpy(),
d.detach().cpu().numpy(),
)

def triple_backward_cpu(
self,
in1,
in2,
out_grad,
weights,
weights_dgrad,
in1_dgrad,
in2_dgrad,
out_tgrad,
weights_tgrad,
in1_tgrad,
in2_tgrad,
graph,
):
assert self.torch_op

in1_torch = torch.tensor(in1).to("cuda").requires_grad_(True)

@vbharadwaj-bk vbharadwaj-bk Jun 2, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

TODO-someday: I wonder if we can combine all of these derivative functions into one to compact this file.

@asglover asglover Jun 5, 2026

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Do you just mean, turn them from .to("cuda") to (... , dtype=cuda).

If so, I should definitely do this. It's considered good practice to. technically this does something funny which is create the tensor, then move it to the device, so the autograd actually does the opposite, moves it off device at the end. When it's supposed to just be on device I ought to use dtype=cuda. I'll clean this up. good point.

@asglover asglover Jun 5, 2026

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

fixed

in2_torch = torch.tensor(in2).to("cuda").requires_grad_(True)
weights_torch = torch.tensor(weights).to("cuda").requires_grad_(True)
out_grad_torch = torch.tensor(out_grad).to("cuda").requires_grad_(True)
in1_dgrad_torch = torch.tensor(in1_dgrad).to("cuda").requires_grad_(True)
in2_dgrad_torch = torch.tensor(in2_dgrad).to("cuda").requires_grad_(True)
weights_dgrad_torch = (
torch.tensor(weights_dgrad).to("cuda").requires_grad_(True)
)
out_tgrad_torch = torch.tensor(out_tgrad).to("cuda")
in1_tgrad_torch = torch.tensor(in1_tgrad).to("cuda")
in2_tgrad_torch = torch.tensor(in2_tgrad).to("cuda")
weights_tgrad_torch = torch.tensor(weights_tgrad).to("cuda")

torch_rows = torch.tensor(graph.rows, device="cuda")
torch_cols = torch.tensor(graph.cols, device="cuda")
torch_transpose_perm = torch.tensor(graph.transpose_perm, device="cuda")

out_torch = self.forward(
in1_torch,
in2_torch,
weights_torch,
torch_rows,
torch_cols,
torch_transpose_perm,
)
in1_grad, in2_grad, weights_grad = torch.autograd.grad(
outputs=out_torch,
inputs=[in1_torch, in2_torch, weights_torch],
grad_outputs=out_grad_torch,
create_graph=True,
retain_graph=True,
)
double_grads = torch.autograd.grad(
outputs=[in1_grad, in2_grad, weights_grad],
inputs=[in1_torch, in2_torch, weights_torch, out_grad_torch],
grad_outputs=[in1_dgrad_torch, in2_dgrad_torch, weights_dgrad_torch],
create_graph=True,
retain_graph=True,
allow_unused=True,
)
double_grads = _none_to_zeros(
double_grads, (in1_torch, in2_torch, weights_torch, out_grad_torch)
)
triple_grads = torch.autograd.grad(
outputs=double_grads,
inputs=[
in1_torch,
in2_torch,
weights_torch,
out_grad_torch,
in1_dgrad_torch,
in2_dgrad_torch,
weights_dgrad_torch,
],
grad_outputs=[
in1_tgrad_torch,
in2_tgrad_torch,
weights_tgrad_torch,
out_tgrad_torch,
],
allow_unused=True,
)
triple_grads = _none_to_zeros(
triple_grads,
(
in1_torch,
in2_torch,
weights_torch,
out_grad_torch,
in1_dgrad_torch,
in2_dgrad_torch,
weights_dgrad_torch,
),
)

return tuple(grad.detach().cpu().numpy() for grad in triple_grads)
123 changes: 122 additions & 1 deletion openequivariance/openequivariance/_torch/TensorProduct.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -200,6 +200,12 @@ def fake_double_backward(kernel, hash, L1_in, L2_in, W, L3_grad, E, F, G):

def register_autograd():
backward_op = torch.ops.libtorch_tp_jit.jit_tp_backward
double_backward_op = torch.ops.libtorch_tp_jit.jit_tp_double_backward

def zero_if_none(grad_output, like):
if grad_output is None:
return torch.zeros_like(like)
return grad_output

def setup_context(ctx, inputs, output):
ctx.kernel, ctx.hash, ctx.L1_in, ctx.L2_in, ctx.weights, ctx.L3_dim = inputs
Expand All @@ -218,7 +224,7 @@ def setup_context_double_backward(ctx, inputs, output):
ctx.kernel, ctx.hash, ctx.L1_in, ctx.L2_in, ctx.weights, ctx.L3_grad = inputs

def double_backward(ctx, E, F, G):
result = torch.ops.libtorch_tp_jit.jit_tp_double_backward(
result = double_backward_op(
ctx.kernel,
ctx.hash,
ctx.L1_in,
Expand All @@ -237,6 +243,121 @@ def double_backward(ctx, E, F, G):
setup_context=setup_context_double_backward,
)

def setup_context_triple_backward(ctx, inputs, output):
(
ctx.kernel,
ctx.hash,
ctx.L1_in,
ctx.L2_in,
ctx.weights,
ctx.L3_grad,
ctx.L1_dgrad,
ctx.L2_dgrad,
ctx.W_dgrad,
) = inputs

def triple_backward(ctx, t_L1_grad, t_L2_grad, t_W_grad, t_L3_dgrad):
t_L1_grad = zero_if_none(t_L1_grad, ctx.L1_in)
t_L2_grad = zero_if_none(t_L2_grad, ctx.L2_in)
t_W_grad = zero_if_none(t_W_grad, ctx.weights)
t_L3_dgrad = zero_if_none(t_L3_dgrad, ctx.L3_grad)

g1_L1_dgrad, g1_L2_dgrad, g1_W, g1_L3_grad = double_backward_op(
ctx.kernel,
ctx.hash,
ctx.L1_dgrad,
ctx.L2_dgrad,
ctx.weights,
ctx.L3_grad,
t_L1_grad,
t_L2_grad,
torch.zeros_like(ctx.weights),
)
g2_L1, g2_L2, g2_W_dgrad, g2_L3_grad = double_backward_op(
ctx.kernel,
ctx.hash,
ctx.L1_in,
ctx.L2_in,
ctx.W_dgrad,
ctx.L3_grad,
t_L1_grad,
t_L2_grad,
torch.zeros_like(ctx.W_dgrad),
)
g3_L1_dgrad, g3_L2, g3_W, g3_L3_grad = double_backward_op(
ctx.kernel,
ctx.hash,
ctx.L1_dgrad,
ctx.L2_in,
ctx.weights,
ctx.L3_grad,
torch.zeros_like(ctx.L1_dgrad),
torch.zeros_like(ctx.L2_in),
t_W_grad,
)
g4_L1, g4_L2_dgrad, g4_W, g4_L3_grad = double_backward_op(
ctx.kernel,
ctx.hash,
ctx.L1_in,
ctx.L2_dgrad,
ctx.weights,
ctx.L3_grad,
torch.zeros_like(ctx.L1_in),
torch.zeros_like(ctx.L2_dgrad),
t_W_grad,
)

g5_L1_dgrad, g5_L2, g5_W = backward_op(
ctx.kernel,
ctx.hash,
ctx.L1_dgrad,
ctx.L2_in,
ctx.weights,
t_L3_dgrad,
)
g6_L1, g6_L2_dgrad, g6_W = backward_op(
ctx.kernel,
ctx.hash,
ctx.L1_in,
ctx.L2_dgrad,
ctx.weights,
t_L3_dgrad,
)
g7_L1, g7_L2, g7_W_dgrad = backward_op(
ctx.kernel,
ctx.hash,
ctx.L1_in,
ctx.L2_in,
ctx.W_dgrad,
t_L3_dgrad,
)

grad_L1 = g2_L1 + g4_L1 + g6_L1 + g7_L1
grad_L2 = g2_L2 + g3_L2 + g5_L2 + g7_L2
grad_W = g1_W + g3_W + g4_W + g5_W + g6_W
grad_L3_grad = g1_L3_grad + g2_L3_grad + g3_L3_grad + g4_L3_grad
grad_L1_dgrad = g1_L1_dgrad + g3_L1_dgrad + g5_L1_dgrad
grad_L2_dgrad = g1_L2_dgrad + g4_L2_dgrad + g6_L2_dgrad
grad_W_dgrad = g2_W_dgrad + g7_W_dgrad

return (
None,
None,
grad_L1,
grad_L2,
grad_W,
grad_L3_grad,
grad_L1_dgrad,
grad_L2_dgrad,
grad_W_dgrad,
)

torch.library.register_autograd(
"libtorch_tp_jit::jit_tp_double_backward",
triple_backward,
setup_context=setup_context_triple_backward,
)


def register_autocast():
global torch
Expand Down
Loading
Loading