code: def setup(app):
      app.add_directive("flax_module", FlaxModuleDirective)
  
      return {
          "version": sphinx.__display_version__,
          "parallel_read_safe": True,
          "parallel_write_safe": True,
      }
  
file_path: flax/docs/_ext/flax_module.py
imports: import importlib
  import sphinx
  import sphinx.ext.autosummary.generate
  from docutils import nodes
  from docutils.parsers.rst import directives
  from docutils.statemachine import ViewList
  from sphinx.util.docutils import SphinxDirective
  from docs.conf_sphinx_patch import generate_autosummary_content
label_desc: Use image node_id 2 for calling a global function with example usage: setup(app) and returns: dict