code: def _module_update_wrapper(
      wrapper: Module,
      wrapped: Optional[Callable[_P, _T]],
      inplace: bool,
  ) -> Callable[_P, _T]:
      cls = wrapper.__class__
      if not isinstance(getattr(cls, "__wrapped__", None), property):
          raise ValueError(
              "Wrapper module must supply `__wrapped__` as a property."
          )
  
      if wrapped is None:
          wrapped = wrapper.__wrapped__  # pyright: ignore
  
      if not inplace:
          # Make a clone, to avoid mutating the original input.
          leaves, treedef = jtu.tree_flatten(wrapper)
          wrapper = jtu.tree_unflatten(treedef, leaves)
  
      initable_cls = _make_initable(cls, None, None, wraps=True)
      object.__setattr__(wrapper, "__class__", initable_cls)
      try:
          # Like `ft.update_wrapper(wrapper, wrapped, updated=())`.
          # We don't update __dict__ as it's common/possible for wrapper and wrapped to
          # both be classes implementing __call__, in which case copying __dict__ over
          # basically just breaks the wrapper class.
          # We don't set __wrapped__, and instead demand that the wrapper class tell us
          # how to redirect to the wrapped object. This is avoid duplicating part of the
          # PyTree.
          for attr in _wrapper_field_names:
              try:
                  value = getattr(wrapped, attr)
              except AttributeError:
                  pass
              else:
                  setattr(wrapper, attr, value)
      finally:
          object.__setattr__(wrapper, "__class__", cls)
      return cast(Callable[_P, _T], wrapper)
  
imports: import abc
  import dataclasses
  import functools
  import inspect
  import types
  import warnings
  import weakref
  from collections.abc import Callable
  from typing import Any, cast, Optional, Protocol, TYPE_CHECKING, TypeVar, Union
  from typing_extensions import dataclass_transform, ParamSpec
  import jax._src.traceback_util
  import jax.tree_util
  import numpy
  from jaxtyping import Array, Bool, PyTreeDef
  from ._better_abstract import ABCMeta, dataclass
  from ._caches import internal_lru_caches
  from ._doc_utils import doc_repr
  from ._pretty_print import tree_pformat
  from ._tree import tree_equal
label_desc: Use image node_id 11 for calling a global function with example usage: _module_update_wrapper(wrapper, wrapped, inplace) and returns: cast