code: def _memoize_by_id(f):
      """Memoizes the result of a higher order function on input function id."""
      cache = {}
  
      @functools.wraps(f)
      def wrapper(g):
          i = id(g)
          try:
              res = cache[i]
          except KeyError:
              res = cache[i] = f(g)
          return res
  
      return wrapper
  
imports: import collections
  import collections.abc
  from collections.abc import Mapping, MutableMapping
  import functools
  import inspect
  from typing import Any, Callable, Optional, TypeVar
  from haiku._src import base
  import jax
  import jax.numpy
label_desc: Use image node_id 15 for calling a global function with example usage: _memoize_by_id(f) and returns: wrapper, res