code: def project_perturbation_pyramid_inf(
      aug_params, epsilon, input_image, image_bounds
  ):
      """Project `perturbation` onto L-infinity ball of radius `epsilon`."""
      del input_image, image_bounds
  
      if epsilon is None:
          return aug_params
  
      # The idea is to ensure that the sum of the perturbations over the pyramid
      # levels can't be more than epsilon.
      clipped_perturbation_pyramid = jax.tree_util.tree_map(
          functools.partial(jnp.clip, a_min=-epsilon, a_max=epsilon),
          aug_params,
      )
  
      return clipped_perturbation_pyramid
  
imports: import functools
  import jax
  import jax.numpy
  import optax
label_desc: Use image node_id 2 for calling a global function with example usage: project_perturbation_pyramid_inf(aug_params, epsilon, input_image, image_bounds) and returns: clipped_perturbation_pyramid, aug_params