code: def annotate(
      f: WrappedFun, in_type: core.InputType | None
  ) -> WrappedFun:
      assert f.in_type is None
      if in_type is None:
          return f
      _check_input_type(in_type)
      return WrappedFun(
          f.f, f.transforms, f.stores, f.params, in_type, f.debug_info
      )
  
imports: from __future__ import annotations
  from functools import partial
  import operator
  from typing import Any, Callable, NamedTuple
  import weakref
  from jax._src import config
  from jax._src import core
  from jax._src import traceback_util
  from jax._src.tree_util import tree_map
  from jax._src.util import curry
label_desc: Use image node_id 5 for calling a global function with example usage: annotate(f, in_type) and returns: WrappedFun, f