code: def tree_map_(
      func: Callable[..., Any],
      tree: PyTree,
      *rests: PyTree,
      is_leaf: Optional[Callable[[PyTree], bool]] = None,
  ) -> PyTree:
      """Like :func:`tree_map`, but do an inplace call on each leaf and return the original tree.
  
      See also :func:`tree_map`.
  
      Args:
          func (callable): A function that takes ``1 + len(rests)`` arguments, to be applied at the
              corresponding leaves of the pytrees.
          tree (pytree): A pytree to be mapped over, with each leaf providing the first positional
              argument to function ``func``.
          rests (tuple of pytree): A tuple of pytrees, each of which has the same structure as
              ``tree`` or has ``tree`` as a prefix.
          is_leaf (callable, optional): An extra leaf predicate function that will be called at each
              flattening step. The function should have a single argument with signature
              ``is_leaf(node) -> bool``. If it returns :data:`True`, the whole subtree being treated
              as a leaf. Otherwise, the default pytree registry will be used to determine a node is a
              leaf or not. If the function is not specified, the default pytree registry will be used.
  
      Returns:
          The original ``tree`` with the value at each leaf is given by the side-effect of function
          ``func(x, *xs)`` (not the return value) where ``x`` is the value at the corresponding leaf
          in ``tree`` and ``xs`` is the tuple of values at values at corresponding nodes in ``rests``.
      """
      return optree.tree_map_(
          func,
          tree,
          *rests,
          is_leaf=is_leaf,
          none_is_leaf=True,
          namespace="torch",
      )
  
imports: import functools
  import warnings
  from typing import Any, Callable, Iterable, List, Optional, overload, Tuple, Type, TypeVar, Union
  import torch
  import optree
  from optree import PyTreeSpec
label_desc: Use image node_id 10 for calling a global function with example usage: tree_map_(func, tree) and returns: optree