Tentative solution to pytree

tests/test_core.py

@@ -1,5 +1,7 @@
+from dataclasses import dataclass
 from functools import partial
 
+import jax
 import jax.numpy as jnp
 import pytest
 import torch
@@ -55,6 +57,36 @@ def test_ones_like():
   t2j_function_test(torch.ones_like, [(2, 3)], tests=tests)
 
 
+def test_tree_coerce():
+  @dataclass
+  class A:
+    a: torch.Tensor
+
+  @dataclass
+  class B:
+    a: A
+    b: int
+
+  torch.utils._pytree.register_dataclass(A)
+  torch.utils._pytree.register_dataclass(B)
+
+  def f():
+    t1 = torch.return_types.max([1, torch.ones(3)])
+    t2 = torch.return_types.topk([3.0, t1])
+    t3 = A(torch.ones(3, 2))
+    t4 = B(t3, 7)
+    return (t2, t3, t4)
+
+  # test jit version works first
+  jax.jit(t2j(f))()
+  # test non-jit version
+  t2j(f)()
+
+  # test the values in the leaves are correct
+  tests = [forward_test]
+  t2j_function_test(f, [], tests=tests)
+
+
 def test_tensor():
   tests = [forward_test]
   t2j_function_test(lambda: torch.tensor([]), [], tests=tests)

torch2jax/__init__.py

@@ -1,5 +1,6 @@
 import copy
 import functools
+from collections import deque
 from contextlib import contextmanager
 from typing import Literal, Optional, Sequence, Tuple, Union
 
@@ -11,6 +12,7 @@
 from torch.overrides import TorchFunctionMode, resolve_name
 from torch.utils._pytree import register_pytree_node as torch_register_pytree_node
 from torch.utils._pytree import tree_map as torch_tree_map
+from torch.utils._pytree import tree_structure as torch_tree_structure
 
 # so that __getitem__ & __setitem__ with mixed keys of int / tensor could work
 torch_register_pytree_node(slice, lambda s: ((s.start, s.stop, s.step), None), lambda values, ctx: slice(*values))
@@ -211,6 +213,33 @@ def _coerce(x):
     )
 
 
+def _tree_coerce(x):
+  spec = torch_tree_structure(x)
+  specs = deque([spec])
+
+  while len(specs) > 0:
+    spec = specs.popleft()
+    specs.extend(spec.children_specs)
+    if spec.is_leaf():
+      continue
+    if spec.type in jax._src.tree_util._registry:
+      continue  # already registered
+    node = torch.utils._pytree.SUPPORTED_NODES.get(spec.type, None)
+    if node is None:
+      continue
+
+    def flip_args(fn):
+      return lambda a, b: fn(b, a)
+
+    jax.tree_util.register_pytree_node(
+      spec.type,
+      node.flatten_fn,
+      flip_args(node.unflatten_fn),  # torch & jax have different order of args
+    )
+
+  return jax.tree.map(_coerce, x)
+
+
 def _v(x):
   assert isinstance(x, Torchish)
   return x.value
@@ -1185,7 +1214,7 @@ def f_jax(*args, rng=None):
         with TorchishMode():
           out = f(*torch_args)
     # use the torch's tree_map, because out is generated from torch code
-    return torch.utils._pytree.tree_map(lambda x: x.value if isinstance(x, Torchish) else x, out)
+    return _tree_coerce(out)
 
   return f_jax