code: def difference(
      before: InternalState, after: InternalState
  ) -> InternalState:
      """Returns an InternalState object with unchanged items set to ``None``.
  
      Note that to determine what values have changed we compare them by identity
      not by value. This is only reasonable to do if `difference` is used to compare
      state *inside* a JAX transform (e.g. comparing the arguments passed into JIT
      with the values that you are about to return from it).
  
      This function never produces false negatives (e.g. we will never incorrectly
      say that a piece of state is unchanged when it has), however it may produce
      false positives. One well known case is if a value is traced by an inner JAX
      transform but unchanged, the identity of the Python object will differ from
      the value passed into the outer function, but the value will not have changed.
      In this case `difference` will say that the value has changed. For example if
      the following change happened inside a function whose state was being diffed
      we would defensively say that ``u`` had changed value even though it had only
      changed Python identity:
  
      >>> u = hk.get_state("u", [], init=jnp.ones)
      >>> u, _ = jax.jit(lambda a: a, a ** 2)(u)
      >>> hk.set_state("u", u)
  
      Args:
        before: state before.
        after: state after.
  
      Returns:
        The difference between before and after, with any values that have the same
        identity before and after set to `None`.
      """
  
      def if_changed(is_new, box_a, box_b):
          if box_a.value is None or is_new(box_a.value, box_b.value):
              return box_b.value
          else:
              return None
  
      # params
      is_new_param = lambda a, b: a is not b
      params_before, params_after = box_and_fill_missing(
          before.params, after.params
      )
      params_after = jax.tree_util.tree_map(
          functools.partial(if_changed, is_new_param),
          params_before,
          params_after,
      )
  
      # state
      def is_new_state(a: base.StatePair, b: base.StatePair):
          return (
              a.initial is not b.initial or a.current is not b.current
          )
  
      state_before, state_after = box_and_fill_missing(
          before.state, after.state
      )
      state_after = jax.tree_util.tree_map(
          functools.partial(if_changed, is_new_state),
          state_before,
          state_after,
      )
  
      # rng
      def is_new_rng(
          a: Optional[base.PRNGSequenceState],
          b: Optional[base.PRNGSequenceState],
      ):
          if a is None:
              return True
          assert len(a) == 2 and len(b) == 2
          return a[0] is not b[0] or a[1] is not b[1]
  
      rng = after.rng if is_new_rng(before.rng, after.rng) else None
  
      return InternalState(params_after, state_after, rng)
  
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 10 for calling a global function with example usage: difference(before, after) and returns: InternalState, unknown, unknown, box_b, None, True