code: def preprocess_inputs_with_augmentation(
      x, data_augmentation=False, rng=None
  ):
      """Input normalization and if specified, data augmentation."""
      out = x.astype(jnp.float32) / 255.0
      if data_augmentation:
          if rng is None:
              raise ValueError("Pass rng when using data augmentation")
          out = drq_image_augmentation(rng, out)
      return out
  
imports: import copy
  import functools
  from absl import logging
  from dopamine.jax import networks
  from dopamine.jax.agents.full_rainbow import full_rainbow_agent
  import gin
  import jax
  import jax.numpy
  import numpy
  import tensorflow
label_desc: Use image node_id 6 for calling a global function with example usage: preprocess_inputs_with_augmentation(x, data_augmentation, rng) and returns: out