code: def get_adv_pyramid(
      batch,
      clean_logits,
      attack_fn,
      optimizer_def,
      config,
      dropout_rng,
      lr_rng,
      misc_artifacts,
  ):
      """Get pyramid attack for batch."""
      del clean_logits, dropout_rng
  
      epsilon = config.advprop.epsilon
      num_steps = config.advprop.num_steps
      # 1 is global, 224 is pixel-level
      # pylint: disable=eval-used
      pyramid_sizes = eval(config.advprop.pyramid_sizes)
      pyramid_scalars = eval(config.advprop.pyramid_scalars)
      if not typecheck(pyramid_sizes):
          pyramid_sizes = [pyramid_sizes]
      if not typecheck(pyramid_scalars):
          pyramid_scalars = [pyramid_scalars]
  
      local_batch_size = batch["inputs"].shape[0]
      init_mode = config.advprop.get("init_mode", default=None)
      init_aug_params = {}
      if init_mode == "zero":
          for layer, pyramid_size in enumerate(pyramid_sizes):
              init_aug_params[str(layer)] = jnp.zeros(
                  dtype=jnp.float32,
                  shape=(
                      local_batch_size,
                      pyramid_size,
                      pyramid_size,
                      3,
                  ),
              )
      elif init_mode == "random":
          for layer, pyramid_size in enumerate(pyramid_sizes):
              init_aug_params[str(layer)] = jax.random.uniform(
                  key=lr_rng,
                  dtype=jnp.float32,
                  shape=(
                      local_batch_size,
                      pyramid_size,
                      pyramid_size,
                      3,
                  ),
                  maxval=epsilon,
                  minval=-epsilon,
              )
      elif init_mode == "normal":
          for layer, pyramid_size in enumerate(pyramid_sizes):
              init_aug_params[str(layer)] = epsilon * jax.random.normal(
                  key=lr_rng,
                  dtype=jnp.float32,
                  shape=(
                      local_batch_size,
                      pyramid_size,
                      pyramid_size,
                      3,
                  ),
              )
      else:
          raise NotImplementedError(
              "Not valid init_mode %s" % init_mode
          )
  
      init_aug_scalars = {}
      if pyramid_scalars is None:
          for layer, pyramid_size in enumerate(pyramid_sizes):
              init_aug_scalars[str(layer)] = jnp.ones(
                  (len(pyramid_sizes))
              )
      else:
          for layer, pyramid_size in enumerate(pyramid_sizes):
              init_aug_scalars[str(layer)] = pyramid_scalars[
                  layer
              ] * jnp.ones((len(pyramid_sizes),))
  
      transform_fn = functools.partial(
          attack_transforms.patched_color_jitter,
          aug_fn=attack_transforms.fast_color_perturb,
      )
  
      def transform_fn_pyramid(input_image, aug_params_dict):
          for layer in range(len(aug_params_dict)):
              input_image = transform_fn(
                  input_image,
                  init_aug_scalars[str(layer)]
                  * aug_params_dict[str(layer)],
              )
          return input_image
  
      (
          adv_image,
          adv_perturbation,
          _,
          attack_artifacts,
      ) = attack_methods.pgd_attack_transform(
          loss_fn=attack_fn,
          transform_fn=transform_fn_pyramid,
          init_aug_params=init_aug_params,
          input_image=batch["inputs"],
          label=batch["label"],
          epsilon=epsilon,
          num_steps=num_steps,
          rng=lr_rng,
          optimizer_def=optimizer_def,
          projection=attack_methods.project_perturbation_pyramid_inf,
      )
      local_result_advprop_pyramid = adv_image, adv_perturbation
      for key in attack_artifacts:
          misc_artifacts[key] = attack_artifacts[key]
  
      return (lambda _: local_result_advprop_pyramid), misc_artifacts
  
file_path: scenic/scenic/projects/adversarialtraining/attacks/attack_compute.py
imports: import functools
  import jax
  import jax.numpy
  import ml_collections
  from scenic.projects.adversarialtraining.attacks import attack_losses
  from scenic.projects.adversarialtraining.attacks import attack_methods
  from scenic.projects.adversarialtraining.attacks import attack_transforms
  from scenic.train_lib import lr_schedules
  from scenic.train_lib import optimizers
label_desc: Use image node_id 2 for calling a global function with example usage: get_adv_pyramid(batch, clean_logits, attack_fn, optimizer_def, config, dropout_rng, lr_rng, misc_artifacts) and returns: function,  misc_artifacts, input_image