code: def flattened_fun_in_tree(
      fn: lu.WrappedFun,
  ) -> tuple[PyTreeDef, Callable[[], PyTreeDef], bool] | None:
      # This implementation relies on internal details of linear_util.py's
      # WrappedFun, but it's for the worthy cause of better user error messages.
      # It can fail (i.e. return None) if its WrappedFun argument is not transformed
      # with flatten_fun or flatten_fun_nokwargs, which could happen e.g. when
      # core.eval_jaxpr encounters a call primitive (though at that point we're just
      # round-tripping jaxprs and the user errors in question are impossible).
      assert (
          isinstance(flatten_fun, partial)
          and len(flatten_fun.args) == 1
      )
      assert (
          isinstance(flatten_fun_nokwargs, partial)
          and len(flatten_fun_nokwargs.args) == 1
      )
      flattens = {flatten_fun.args[0], flatten_fun_nokwargs.args[0]}
      try:
          (((in_tree,), out_tree_store, has_kwargs),) = (
              (args, store, f is flatten_fun.args[0])
              for (f, args), store in zip(fn.transforms, fn.stores)
              if f in flattens
          )
      except ValueError:
          return None
      else:
          return in_tree, lambda: out_tree_store.val, has_kwargs
  
imports: from __future__ import annotations
  from collections.abc import Iterable, Sequence
  import inspect
  import operator
  from functools import partial
  from typing import Any, Callable
  import warnings
  import numpy
  from jax._src import core
  from jax._src import dtypes
  from jax._src.abstract_arrays import numpy_scalar_types
  from jax._src.core import ShapedArray
  from jax._src.tree_util import PyTreeDef, tree_flatten, tree_unflatten, tree_map, tree_structure, treedef_children, generate_key_paths, keystr, broadcast_prefix, prefix_errors
  from jax._src.tree_util import _replace_nones
  from jax._src import linear_util
  from jax._src.linear_util import TracingDebugInfo
  from jax._src.util import safe_map, WrapKwArgs, Hashable, HashableFunction, Unhashable
  from jax._src import traceback_util
label_desc: Use image node_id 9 for calling a global function with example usage: flattened_fun_in_tree(fn) and returns: in_tree,  function,  has_kwargs, None