code: def build_attack_fn(
      model, transform_name, attack_name, use_logits=True
  ):
      """Builds a compiled tf.function for given transform and attack.
  
      Args:
        model: A callable which makes inference on input.
        transform_name: The name of the transform to be used before the attack. See
          transforms.TRANSFORM_FUNCTIONS.
        attack_name: The name of the attack to be used. See
          attacks.ATTACK_CONSTRUCTORS.
        use_logits: A Boolean indicating whether the model returns logits or class
          probabilities.
  
      Returns:
        A compiled tf.function which generates adversarial input on given input and
        label.
      """
      transform_fn = transforms.TRANSFORM_FUNCTIONS[transform_name]
      inv_transform_fn = transforms.INVERSE_TRANSFORM_FUNCTIONS[
          transform_name
      ]
      attack = attacks.construct_attack(attack_name)
  
      pgd_model_fn = lambda x: model(inv_transform_fn(x))
      pgd_loss_fn = tf.keras.losses.SparseCategoricalCrossentropy(
          from_logits=use_logits,
          reduction=tf.keras.losses.Reduction.NONE,
      )
  
      @tf.function
      def attack_fn(x, y):
          # transform -> attack -> inverse transform -> clip.
          transformed_x = transform_fn(x)
          adv_x = attack.attack(
              transformed_x,
              y,
              pgd_model_fn,
              pgd_loss_fn,
              random_start=True,
          )
          return tf.clip_by_value(inv_transform_fn(adv_x), 0.0, 1.0)
  
      return attack_fn
  
file_path: neural-structured-learning/research/multi_representation_adversary/multi_representation_adversary/helper.py
imports: from absl import logging
  from multi_representation_adversary import attacks
  from multi_representation_adversary import transforms
  import tensorflow.compat.v2
label_desc: Use image node_id 1 for calling a global function with example usage: build_attack_fn(model, transform_name, attack_name, use_logits) and returns: attack_fn, tf