code: def tree_deserialise_leaves(
      path_or_file: Union[str, pathlib.Path, BinaryIO],
      like: PyTree,
      filter_spec=default_deserialise_filter_spec,
      is_leaf: Optional[Callable[[Any], bool]] = None,
  ) -> PyTree:
      """Load the leaves of a PyTree from a file.
  
      **Arguments:**
  
      - `path_or_file`: The file location to load values from or a binary file-like
          object.
      - `like`: A PyTree of same structure, and with leaves of the same type, as the
          PyTree being loaded. Those leaves which are loaded will replace the
          corresponding leaves of `like`.
      - `filter_spec`: Specifies how to load each kind of leaf. By default all JAX
          arrays, NumPy arrays, Python bool/int/float/complexes are loaded, and
          all other leaf types are not loaded, and will retain their
          value from `like`. (See [`equinox.default_deserialise_filter_spec`][].)
      - `is_leaf`: Called on every node of `like`; if `True` then this node will be
          treated as a leaf.
  
      **Returns:**
  
      The loaded PyTree, formed by iterating over `like` and replacing some of its leaves
      with the leaves saved in `path`.
  
      !!! example
  
          This can be used to load a model from file.
  
          ```python
          import equinox as eqx
          import jax.random as jr
  
          model_original = eqx.nn.MLP(2, 2, 2, 2, key=jr.PRNGKey(0))
          eqx.tree_serialise_leaves("some_filename.eqx", model_original)
          model_loaded = eqx.tree_deserialise_leaves("some_filename.eqx", model_original)
  
          # To partially load weights, do model surgery. In this case load everything
          # except the final layer.
          model_partial = eqx.tree_at(lambda mlp: mlp.layers[-1], model_loaded, model_original)
          ```
  
      !!! example
  
          A common pattern is the following:
  
          ```python
          def run(..., load_path=None):
              if load_path is None:
                  model = Model(...hyperparameters...)
              else:
                  model = eqx.filter_eval_shape(Model, ...hyperparameters...)
                  model = eqx.tree_deserialise_leaves(load_path, model)
          ```
          in which either a model is created directly (e.g. at the start of training), or
          a suitable `like` is constructed (e.g. when resuming training), where
          [`equinox.filter_eval_shape`][] is used to avoid creating spurious short-lived
          arrays taking up memory.
  
      !!! info
  
          `filter_spec` should typically be a function `(File, Any) -> Any`, which takes
          a file handle and a leaf from `like`, and either returns the corresponding
          loaded leaf, or retuns the leaf from `like` unchanged.
  
          It can also be a PyTree of such functions, in which case the PyTree structure
          should be a prefix of `pytree`, and each function will be mapped over the
          corresponding sub-PyTree of `pytree`.
      """  # noqa: E501
      with _maybe_open(path_or_file, "rb") as f:
  
          def _deserialise(spec, x):
              def __deserialise(y):
                  return spec(f, y)
  
              return _ordered_tree_map(
                  __deserialise, x, is_leaf=is_leaf
              )
  
          out = _ordered_tree_map(_deserialise, filter_spec, like)
      with jax.ensure_compile_time_eval():
          # ArrayImpl isn't a public type, so this is how we get access to it instead.
          # `ensure_compile_time_eval` just in case someone is doing deserialisation
          # inside JIT. Which would be weird, but still.
          array_impl_type = type(jnp.array(0))
      jtu.tree_map_with_path(
          _assert_same(array_impl_type), out, like, is_leaf=is_leaf
      )
      return out
  
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 8 for calling a global function with example usage: tree_deserialise_leaves(path_or_file, like, filter_spec, is_leaf) and returns: out, _ordered_tree_map, spec