code: def transform_and_run_once(f, *args, **kwargs):
      f = transform.transform(f)
  
      def g(*args, **kwargs):
          rng = jax.random.PRNGKey(28)
          params = f.init(rng, *args, **kwargs)
          out = f.apply(params, None, *args, **kwargs)
          return params, out
  
      return jax.tree_util.tree_map(
          lambda x: x.dtype, jax.eval_shape(g, *args, **kwargs)
      )
  
imports: import importlib
  from typing import Optional
  from absl.testing import absltest
  from haiku._src import base
  from haiku._src import conv
  from haiku._src import mixed_precision
  from haiku._src import module
  from haiku._src import test_utils
  from haiku._src import transform
  import jax
  import jax.numpy
  import jmp
label_desc: Use image node_id 2 for calling a global function with example usage: transform_and_run_once(f) and returns: jax, params,  out