torch.compile friendly pytree

Author: lucidrains

lumiere_pytorch/lumiere.py

@@ -20,8 +20,7 @@
 
 from einops import rearrange, pack, unpack, repeat
 
-from optree import tree_flatten, tree_unflatten
-import torch.utils._pytree as pytree
+from torch.utils._pytree import tree_flatten, tree_unflatten
 
 from x_transformers.x_transformers import (
     Attention,
@@ -547,11 +546,8 @@ def auto_repeat_tensors_for_time(_, args, kwargs):
             first_arg, *rest_args = args
 
             all_rest_args = (rest_args, kwargs)
-            if torch.compiler.is_dynamo_compiling():
-                all_rest_args, pytree_spec = pytree.tree_flatten(all_rest_args)
-            else:
-                all_rest_args, pytree_spec = tree_flatten(all_rest_args)
 
+            all_rest_args, pytree_spec = tree_flatten(all_rest_args)
 
             if not is_tensor(first_arg) or len(all_rest_args) == 0:
                 return args, kwargs
@@ -571,11 +567,8 @@ def auto_repeat_tensors_for_time(_, args, kwargs):
 
                 out_rest_args.append(arg)
 
-            if torch.compiler.is_dynamo_compiling():
-                # reordering of args is deliberate. pytree and optree have different API
-                rest_args, kwargs = pytree.tree_unflatten(out_rest_args, pytree_spec)
-            else:
-                rest_args, kwargs = tree_unflatten(pytree_spec, out_rest_args)
+            rest_args, kwargs = tree_unflatten(out_rest_args, pytree_spec)
+
             return (first_arg, *rest_args), kwargs
 
         for module in self.modules():

setup.py

@@ -3,7 +3,7 @@
 setup(
   name = 'lumiere-pytorch',
   packages = find_packages(exclude=[]),
-  version = '0.0.23',
+  version = '0.0.24',
   license='MIT',
   description = 'Lumiere',
   author = 'Phil Wang',