code: def lazy_init(fn):
      """Lazily evaluates a function by using the shapes of the inputs.
  
      The returned function accepts a combination of JAX values and
      ``jax.ShapeDtypeStruct`` instances for the inputs for which we
      don't need concrete values (only the shape and dtype).
  
      This API is used by ``core.lazy_init`` or ``Module.lazy_init``
      to initialize variables without doing any actual computation on the
      inputs.
  
      Args:
        fn: the function to be lazily evaluated.
      Returns:
        A new function that accepts a mix of concrete values and
        ``jax.ShapeDtypeStruct`` instances.
      """
  
      @functools.wraps(fn)
      def wrapper(*args, **kwargs):
          # TODO(mattjj,jheek): use a public JAX API
          # flatten fn and prepare for internal JAX transform
          inputs_flat, in_tree = jax.tree_util.tree_flatten(
              (args, kwargs)
          )
          f_flat, out_tree = jax.api_util.flatten_fun(
              lu.wrap_init(fn), in_tree
          )
          # map inputs to PartialVal known/unknown
          # only the computations depending on knowns will be executed
          in_pvals = [_maybe_unknown(x) for x in inputs_flat]
          _, out_pvals, _ = pe.trace_to_jaxpr_nounits(f_flat, in_pvals)
          # all outputs should be knowns. If this fails
          # the user is creating variables that depend on a
          # argument that was passed as a ShapeDtypeStruct.
          out_flat = []
          for pv, const in out_pvals:
              if pv is None:
                  # const is the actual value of the known output
                  out_flat.append(const)
              else:
                  raise errors.LazyInitError(pv)
          return jax.tree_util.tree_unflatten(out_tree(), out_flat)
  
      return wrapper
  
file_path: flax/flax/core/partial_eval.py
imports: import functools
  from typing import Any
  import jax
  from jax import core
  from jax.extend import linear_util
  from jax.interpreters import partial_eval
  from flax import errors
label_desc: Use image node_id 2 for calling a global function with example usage: lazy_init(fn) and returns: wrapper, jax