code: def key_for_module(cls: type[hk.Module]) -> ClassInfoOrType:
      """Returns a suitable key for the given module class."""
      if "<locals>" in cls.__qualname__:
          # Some APIs (e.g. `hk.to_module`) are factory functions that create modules.
          # It is not desirable for us to use the qualname in this case since that
          # would associate all class instances created by the factory with a single
          # policy. Instead we use the class object itself, with the assumption that
          # these types are less likely to be created as a side effect of force
          # reloading modules.
          return cls
      else:
          return ClassInfo(cls.__module__, cls.__qualname__)
  
imports: import collections
  import contextlib
  import threading
  from typing import TypeVar, Optional, Union
  from haiku._src import base
  from haiku._src import data_structures
  from haiku._src import module
  import jmp
label_desc: Use image node_id 1 for calling a global function with example usage: key_for_module(cls) and returns: cls, ClassInfo