code: def random_inputs(rng, input_shape):
      if type(input_shape) is tuple:
          return rng.randn(*input_shape).astype(
              dtypes.canonicalize_dtype(float)
          )
      elif type(input_shape) is list:
          return [random_inputs(rng, shape) for shape in input_shape]
      else:
          raise TypeError(type(input_shape))
  
imports: from absl.testing import absltest
  import numpy
  from jax._src import test_util
  from jax import random
  from jax.example_libraries import stax
  from jax import dtypes
  from jax import config
label_desc: Use image node_id 1 for calling a global function with example usage: random_inputs(rng, input_shape) and returns: rng, unknown