code: def flatten(
      inputs: types.ArrayHolder,
  ) -> tp.Iterable[types.np.ndarray]:
      if isinstance(inputs, (jnp.ndarray, np.ndarray)):
          yield inputs
      elif isinstance(inputs, tp.Tuple):
          for x in inputs:
              yield from flatten(x)
      elif isinstance(inputs, tp.Dict):
          for x in inputs.values():
              yield from flatten(x)
      elif isinstance(inputs, (tp.Generator, tp.Iterator, tp.Iterable)):
          yield inputs
      else:
          raise TypeError(f"Unsupported type '{type(inputs)}'")
  
imports: import collections
  import functools
  import math
  import typing
  import jax
  import jax.numpy
  import numpy
  import six
  import typing_extensions
  from elegy import types, utils
label_desc: Use image node_id 3 for calling a global function with example usage: flatten(inputs) without return types