code: def initialized(key, image_size, model):
      input_shape = (1, image_size, image_size, 3)
  
      @jax.jit
      def init(*args):
          return model.init(*args)
  
      variables = init(
          {"params": key}, jnp.ones(input_shape, model.dtype)
      )
      return variables["params"], variables["batch_stats"]
  
imports: import functools
  import time
  from typing import Any
  from absl import logging
  from clu import metric_writers
  from clu import periodic_actions
  from flax import jax_utils
  from flax.training import checkpoints
  from flax.training import common_utils
  from flax.training import dynamic_scale
  from flax.training import train_state
  import jax
  from jax import lax
  import jax.numpy
  from jax import random
  import ml_collections
  import optax
  import tensorflow
  import tensorflow_datasets
  import input_pipeline
  import models
label_desc: Use image node_id 2 for calling a global function with example usage: initialized(key, image_size, model) and returns: variables,  variables, model