code: def only_changed_state(
      old_state: hk.State, new_state: hk.State
  ) -> hk.State:
      """Returns the subset of new_state that has changed from old_state."""
      updated_state = {}
      for mod_name, name, value in hkds.traverse(new_state):
          if mod_name in old_state and name in old_state[mod_name]:
              if old_state[mod_name][name] is not value:
                  if mod_name not in updated_state:
                      updated_state[mod_name] = {}
                  updated_state[mod_name][name] = value
      return updated_state
  
file_path: dm-haiku/haiku/_src/flax/flax_module.py
imports: from typing import TypeVar, Union
  import flax.core
  import flax.linen
  from haiku._src import filtering
  from haiku._src import transform
  from haiku._src import typing
  from haiku._src.flax import utils
label_desc: Use image node_id 2 for calling a global function with example usage: only_changed_state(old_state, new_state) and returns: updated_state