take care of functions that output multiple tensors (say auxiliary lo…

…sses)
lucidrains · Oct 26, 2021 · fdf94c1 · fdf94c1
1 parent c61d8d7
commit fdf94c1
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 3 deletions.
diff --git a/jax2torch/jax2torch.py b/jax2torch/jax2torch.py
@@ -38,8 +38,12 @@ def forward(ctx, *args):
                 return tree_j2t(y_)
 
             @staticmethod
-            def backward(ctx, grad_y):
-                grads, *_ = ctx.fun_vjp(t2j(grad_y))
+            def backward(ctx, *grad_args):
+                if len(grad_args) > 1:
+                    grad_args = tree_t2j(grad_args)
+                else:
+                    grad_args = t2j(grad_args[0])
+                grads, *_ = ctx.fun_vjp(grad_args)
                 ret = tree_j2t(grads), *((None,) * (ctx.num_args - 1))
                 return ret
 

diff --git a/setup.py b/setup.py
@@ -3,7 +3,7 @@
 setup(
   name = 'jax2torch',
   packages = find_packages(exclude=[]),
-  version = '0.0.5',
+  version = '0.0.6',
   license='MIT',
   description = 'Jax 2 Torch',
   author = 'Phil Wang',