code: def arbitrary_mergeable_leaf(min_num_dims, args, kwargs):
      for a in jax.tree_util.tree_leaves(args):
          if ndim_at_least(a, min_num_dims):
              return a
      for k in jax.tree_util.tree_leaves(kwargs):
          if ndim_at_least(k, min_num_dims):
              return k
      # Couldn't find a satisfactory leaf.
      return None
  
imports: from collections.abc import Iterable, Sequence
  import functools
  from typing import Any, Callable, Optional
  from haiku._src import base
  from haiku._src import initializers
  from haiku._src import module
  from haiku._src import typing
  from haiku._src.typing import PRNGKey
  import jax
  from jax import lax
  import jax.numpy
  import numpy
label_desc: Use image node_id 4 for calling a global function with example usage: arbitrary_mergeable_leaf(min_num_dims, args, kwargs) and returns: None, a, k