code: def _ordered_tree_map(
      f: Callable[..., Any],
      tree: Any,
      *rest: Any,
      is_leaf: Optional[Callable[[Any], bool]] = None,
  ) -> Any:
      """Like jax.tree_util.tree_map, but guaranteed to iterate over the tree
      in fixed order. (Namely depth-first left-to-right.)
      """
      # Discussion: https://github.com/patrick-kidger/equinox/issues/136
      paths_and_leaves, treedef = jtu.tree_flatten_with_path(
          tree, is_leaf
      )
      all_leaves = list(zip(*paths_and_leaves)) + [
          treedef.flatten_up_to(r) for r in rest
      ]
  
      @ft.wraps(f)
      def _f(path, *xs):
          try:
              return f(*xs)
          except TreePathError as e:
              combo_path = path + e.path
              exc = TreePathError(
                  f"Error at leaf with path {combo_path}"
              )
              exc.path = combo_path
              raise exc from e
          except Exception as e:
              exc = TreePathError(f"Error at leaf with path {path}")
              exc.path = path
              raise exc from e
  
      return treedef.unflatten(_f(*xs) for xs in zip(*all_leaves))
  
imports: import functools
  import pathlib
  from collections.abc import Callable
  from contextlib import contextmanager
  from typing import Any, BinaryIO, Optional, Union
  import jax
  import jax.numpy
  import jax.tree_util
  import numpy
  from jaxtyping import PyTree
  from ._filters import is_array_like
label_desc: Use image node_id 1 for calling a global function with example usage: _ordered_tree_map(f, tree) and returns: treedef, f