Replace remaining uses of dm-tree with pytrees (#480)

* replace dm-tree with pytree

* fix tests

Author: eb8680

effectful/handlers/jax/_handlers.py

@@ -10,8 +10,6 @@
 except ImportError:
     raise ImportError("JAX is required to use effectful.handlers.jax")
 
-import tree
-
 from effectful.internals.runtime import interpreter
 from effectful.ops.semantics import apply, evaluate, fvsof, typeof
 from effectful.ops.syntax import (
@@ -93,7 +91,7 @@ def _is_eager(t):
 
     if not (
         isinstance(t, Term)
-        and all(_is_eager(a) for a in tree.flatten((t.args, t.kwargs)))
+        and all(_is_eager(a) for a in jax.tree.flatten((t.args, t.kwargs))[0])
     ):
         return t
 
@@ -118,7 +116,7 @@ def reindex_flat_array(t):
         result = jnp.reshape(t, result_shape)
         return jax_getitem(result, tuple(k() for k in sized_fvs.keys()))
 
-    result = tree.map_structure(reindex_flat_array, flat_result)
+    result = jax.tree.map(reindex_flat_array, flat_result)
     return result
 
 
@@ -146,9 +144,9 @@ def _jax_op(*args, **kwargs) -> jax.Array:
             # which partial_eval handles.
             return typing.cast(jax.Array, _partial_eval(tm))
         elif not any(
-            tree.flatten(
-                tree.map_structure(lambda x: isinstance(x, Term), (args, kwargs))
-            )
+            jax.tree.flatten(
+                jax.tree.map(lambda x: isinstance(x, Term), (args, kwargs))
+            )[0]
         ):
             return typing.cast(jax.Array, jax_fn(*args, **kwargs))
         else:
@@ -165,9 +163,9 @@ def _register_jax_op_no_partial_eval[**P, T](jax_fn: Callable[P, T]):
     @defop
     def _jax_op(*args, **kwargs) -> jax.Array:
         if not any(
-            tree.flatten(
-                tree.map_structure(lambda x: isinstance(x, Term), (args, kwargs))
-            )
+            jax.tree.flatten(
+                jax.tree.map(lambda x: isinstance(x, Term), (args, kwargs))
+            )[0]
         ):
             return typing.cast(jax.Array, jax_fn(*args, **kwargs))
         else:
@@ -210,11 +208,9 @@ def bind_dims[T, A, B](
     >>> bind_dims(t, b, a).shape
     (3, 2)
     """
-    if tree.is_nested(value):
-        return tree.map_structure(lambda v: bind_dims(v, *names), value)
-
-    semantic_type = typeof(value)
-    return __dispatch(semantic_type)(value, *names)
+    if jax.tree_util.treedef_is_leaf(jax.tree.structure(value)):
+        return __dispatch(typeof(value))(value, *names)
+    return jax.tree.map(lambda v: bind_dims(v, *names), value)
 
 
 @defop
@@ -225,11 +221,9 @@ def unbind_dims[T, A, B](
     *names: Annotated[Operation[[], jax.Array], Scoped[B]],
 ) -> Annotated[T, Scoped[A | B]]:
     """Convert positional dimensions to named dimensions."""
-    if tree.is_nested(value):
-        return tree.map_structure(lambda v: unbind_dims(v, *names), value)
-
-    semantic_type = typeof(value)
-    return __dispatch(semantic_type)(value, *names)
+    if jax.tree_util.treedef_is_leaf(jax.tree.structure(value)):
+        return __dispatch(typeof(value))(value, *names)
+    return jax.tree.map(lambda v: unbind_dims(v, *names), value)
 
 
 def jit(f, *args, **kwargs):
@@ -244,7 +238,7 @@ def _indexed_func_wrapper[**P, S, T](
     # index expressions for the result of the function
     indexes = None
 
-    # hide index lists from tree.map_structure
+    # hide index lists from jax.tree.mapping
     class Indexes:
         def __init__(self, sizes):
             self.sizes = sizes
@@ -260,21 +254,16 @@ def deindex_tensor(t, i):
             return t_
 
         ret = func(*args, **kwargs)
-        indexes = tree.map_structure(lambda t: Indexes(sizesof(t)), ret)
-        tensors = tree.map_structure(lambda t, i: deindex_tensor(t, i), ret, indexes)
+        indexes = jax.tree.map(lambda t: Indexes(sizesof(t)), ret)
+        tensors = jax.tree.map(lambda t, i: deindex_tensor(t, i), ret, indexes)
         return tensors
 
     # reapply the stored indexes to a result
     def reindex(ret, starting_dim=0):
         def index_expr(i):
             return (slice(None),) * (starting_dim) + tuple(x() for x in i.indexes)
 
-        if tree.is_nested(ret):
-            indexed_ret = tree.map_structure(
-                lambda t, i: getitem(t, index_expr(i)), ret, indexes
-            )
-        else:
-            indexed_ret = getitem(ret, index_expr(indexes))
+        indexed_ret = jax.tree.map(lambda t, i: getitem(t, index_expr(i)), ret, indexes)
 
         return indexed_ret
 

effectful/handlers/jax/_terms.py

@@ -4,7 +4,6 @@
 from typing import Any, cast
 
 import jax
-import tree
 
 import effectful.handlers.jax.numpy as jnp
 from effectful.handlers.jax._handlers import (
@@ -455,10 +454,10 @@ def _bind_dims_array(t: jax.Array, *args: Operation[[], jax.Array]) -> jax.Array
 
     def _evaluate(expr):
         if isinstance(expr, Term):
-            (args, kwargs) = tree.map_structure(_evaluate, (expr.args, expr.kwargs))
+            (args, kwargs) = jax.tree.map(_evaluate, (expr.args, expr.kwargs))
             return _partial_eval(expr)
-        if tree.is_nested(expr):
-            return tree.map_structure(_evaluate, expr)
+        if not jax.tree_util.treedef_is_leaf(jax.tree.structure(expr)):
+            return jax.tree.map(_evaluate, expr)
         return expr
 
     if not isinstance(t, Term):

effectful/handlers/torch.py

@@ -9,7 +9,7 @@
 except ImportError:
     raise ImportError("PyTorch is required to use effectful.handlers.torch")
 
-import tree
+import torch.utils._pytree as pytree
 
 from effectful.internals.runtime import interpreter
 from effectful.internals.tensor_utils import _desugar_tensor_index
@@ -91,7 +91,7 @@ def _partial_eval(t: Expr[torch.Tensor]) -> Expr[torch.Tensor]:
         isinstance(t, Term)
         and all(
             isinstance(a, torch.Tensor) or not isinstance(a, Term) or a.op in sized_fvs
-            for a in tree.flatten((t.args, t.kwargs))
+            for a in pytree.tree_flatten((t.args, t.kwargs))[0]
         )
     ):
         return t
@@ -126,7 +126,7 @@ def reindex_flat_tensor(t):
         result = t.reshape(inds[0].shape + t.shape[1:])
         return torch_getitem(result, tuple(k() for k in sized_fvs.keys()))
 
-    result = tree.map_structure(reindex_flat_tensor, flat_result)
+    result = pytree.tree_map(reindex_flat_tensor, flat_result)
     return result
 
 
@@ -135,9 +135,7 @@ def reindex_flat_tensor(t):
 def bind_dims[
     A,
     B,
-    HasDims: torch.Tensor
-    | torch.distributions.Distribution
-    | tree.Structure[torch.Tensor | torch.distributions.Distribution],
+    HasDims: pytree.PyTree | torch.Tensor | torch.distributions.Distribution,
 ](
     value: Annotated[HasDims, Scoped[A | B]],
     *names: Annotated[Operation[[], torch.Tensor], Scoped[B]],
@@ -157,8 +155,8 @@ def bind_dims[
     >>> bind_dims(t[a(), b()], b, a).shape
     torch.Size([3, 2])
     """
-    if tree.is_nested(value):
-        return tree.map_structure(lambda v: bind_dims(v, *names), value)
+    if not pytree.tree_is_leaf(value):
+        return pytree.tree_map(lambda v: bind_dims(v, *names), value)
     raise NotHandled
 
 
@@ -209,15 +207,13 @@ def _bind_dims_tensor(
 def unbind_dims[
     A,
     B,
-    HasDims: torch.Tensor
-    | torch.distributions.Distribution
-    | tree.Structure[torch.Tensor | torch.distributions.Distribution],
+    HasDims: pytree.PyTree | torch.Tensor | torch.distributions.Distribution,
 ](
     value: Annotated[HasDims, Scoped[A | B]],
     *names: Annotated[Operation[[], torch.Tensor], Scoped[B]],
 ) -> Annotated[HasDims, Scoped[A | B]]:
-    if tree.is_nested(value):
-        return tree.map_structure(lambda v: unbind_dims(v, *names), value)
+    if not pytree.tree_is_leaf(value):
+        return pytree.tree_map(lambda v: unbind_dims(v, *names), value)
     raise NotHandled
 
 
@@ -256,9 +252,9 @@ def _torch_op(*args, **kwargs) -> torch.Tensor:
             # which partial_eval handles.
             return typing.cast(torch.Tensor, _partial_eval(tm))
         elif not any(
-            tree.flatten(
-                tree.map_structure(lambda x: isinstance(x, Term), (args, kwargs))
-            )
+            pytree.tree_flatten(
+                pytree.tree_map(lambda x: isinstance(x, Term), (args, kwargs))
+            )[0]
         ):
             return typing.cast(torch.Tensor, torch_fn(*args, **kwargs))
         else:
@@ -600,7 +596,7 @@ def _indexed_func_wrapper[**P, S, T](
     # index expressions for the result of the function
     indexes = None
 
-    # hide index lists from tree.map_structure
+    # hide index lists from pytree.tree_map
     class Indexes:
         def __init__(self, sizes):
             self.sizes = sizes
@@ -616,17 +612,17 @@ def deindex_tensor(t, i):
             return t_
 
         ret = func(*args, **kwargs)
-        indexes = tree.map_structure(lambda t: Indexes(sizesof(t)), ret)
-        tensors = tree.map_structure(lambda t, i: deindex_tensor(t, i), ret, indexes)
+        indexes = pytree.tree_map(lambda t: Indexes(sizesof(t)), ret)
+        tensors = pytree.tree_map(lambda t, i: deindex_tensor(t, i), ret, indexes)
         return tensors
 
     # reapply the stored indexes to a result
     def reindex(ret, starting_dim=0):
         def index_expr(i):
             return (slice(None),) * (starting_dim) + tuple(x() for x in i.indexes)
 
-        if tree.is_nested(ret):
-            indexed_ret = tree.map_structure(
+        if not pytree.tree_is_leaf(ret):
+            indexed_ret = pytree.tree_map(
                 lambda t, i: torch_getitem(t, index_expr(i)), ret, indexes
             )
         else:
@@ -677,7 +673,7 @@ def jvp(func, *args, **kwargs):
     # hide deindexed_func from _register_torch_op
     jvp_func = functools.partial(torch.func.jvp, deindexed_func)
     ret = _register_torch_op(jvp_func)(*args, **kwargs)
-    return tree.map_structure(reindex, ret)
+    return pytree.tree_map(reindex, ret)
 
 
 @functools.wraps(torch.func.vjp)
@@ -699,15 +695,15 @@ def repack_primals(primals):
     def wrapper(*primals):
         nonlocal indexed_result
         indexed_result = func(*repack_primals(primals))
-        return tree.map_structure(
+        return pytree.tree_map(
             lambda t: bind_dims(t, *list(sizesof(t).keys())), indexed_result
         )
 
     unindexed_primals = [t[0] for t in unpacked_primals]
     _, vjpfunc = torch.func.vjp(wrapper, *unindexed_primals, **kwargs)
 
     def vjpfunc_wrapper(*tangents):
-        unindexed_tangents = tree.map_structure(
+        unindexed_tangents = pytree.tree_map(
             lambda t: bind_dims(t, *list(sizesof(t).keys())), tangents
         )
         grads = vjpfunc(*unindexed_tangents)

pyproject.toml

@@ -36,10 +36,10 @@ Source = "https://github.com/BasisResearch/effectful"
 "Bug Tracker" = "https://github.com/BasisResearch/effectful/issues"
 
 [project.optional-dependencies]
-torch = ["torch", "dm-tree"]
-pyro = ["pyro-ppl>=1.9.1", "dm-tree"]
-jax = ["jax", "dm-tree"]
-numpyro = ["numpyro>=0.19", "dm-tree"]
+torch = ["torch"]
+pyro = ["pyro-ppl>=1.9.1"]
+jax = ["jax"]
+numpyro = ["numpyro>=0.19"]
 llm = [
   "litellm",
   "pillow",