code: def _max_depth(g: Graph) -> int:
      if g.subgraphs:
          return 1 + max(0, *[_max_depth(s) for s in g.subgraphs])
      else:
          return 1
  
imports: import collections
  import contextlib
  import functools
  import html
  from typing import Any, Callable, NamedTuple, Optional
  from haiku._src import data_structures
  from haiku._src import module
  from haiku._src import utils
  import jax
  import jax.core
  from jax.experimental import pjit
  from jax.extend import linear_util
label_desc: Use image node_id 8 for calling a global function with example usage: _max_depth(g) and returns: unknown, int