code: def maximize_within_unit_norm(weights, norm_type, epsilon=1e-6):
      """Solves the maximization problem weights^T*x with the constraint norm(x)=1.
  
      This op solves a batch of maximization problems at one time. The first axis of
      `weights` is assumed to be the batch dimension, and each "row" is treated as
      an independent maximization problem.
  
      This op is mainly used to generate adversarial examples (e.g., FGSM proposed
      by Goodfellow et al.). Specifically, the `weights` are gradients, and `x` is
      the adversarial perturbation. The desired perturbation is the one causing the
      largest loss increase. In this op, the loss increase is approximated by the
      dot product between the gradient and the perturbation, as in the first-order
      Taylor approximation of the loss function.
  
      Args:
        weights: A `Tensor` or a collection of `Tensor` objects representing a batch
          of weights to define the maximization objective. If this is a collection,
          the first dimension of all `Tensor` objects should be the same (i.e. batch
          size).
        norm_type: One of `nsl.configs.NormType`, the type of the norm in the
          constraint.
        epsilon: A lower bound value for the norm to avoid division by 0.
  
      Returns:
        A `Tensor` or a collection of `Tensor` objects (with the same structure and
        shape as `weights`) representing a batch of adversarial perturbations as the
        solution to the maximization problems.
      """
      if isinstance(norm_type, six.string_types):
          # Allows string to be converted into NormType.
          norm_type = configs.NormType(norm_type)
  
      if norm_type == configs.NormType.INFINITY:
          return tf.nest.map_structure(tf.sign, weights)
  
      tensors = tf.nest.flatten(weights)
      tensor_ranks = [t.shape.rank for t in tensors]
  
      if not tensors:  # `weights` is an empty collection.
          return weights
  
      if norm_type == configs.NormType.L2:
          squared_norm = _reduce_across_tensors(
              tf.reduce_sum, [tf.square(t) for t in tensors]
          )
          inv_global_norm = tf.math.rsqrt(
              tf.maximum(squared_norm, epsilon**2)
          )
          normalized_tensors = [
              tensor * _expand_to_rank(inv_global_norm, rank)
              for tensor, rank in zip(tensors, tensor_ranks)
          ]
          return tf.nest.pack_sequence_as(weights, normalized_tensors)
      elif norm_type == configs.NormType.L1:
          # For L1 norm, the solution is to put 1 or -1 at a dimension with maximum
          # absolute value, and 0 at others. In case of multiple dimensions having the
          # same maximum absolute value, any distribution among them will do. Here we
          # choose to distribute evenly among those dimensions for efficient
          # implementation.
          abs_tensors = [tf.abs(t) for t in tensors]
          max_elem = _reduce_across_tensors(tf.reduce_max, abs_tensors)
          is_max_elem = [
              tf.cast(
                  tf.equal(t, _expand_to_rank(max_elem, rank)), t.dtype
              )
              for t, rank in zip(abs_tensors, tensor_ranks)
          ]
          num_nonzero = _reduce_across_tensors(
              tf.reduce_sum, is_max_elem
          )
          denominator = tf.maximum(num_nonzero, epsilon)
          mask = [
              is_max * tf.sign(t) / _expand_to_rank(denominator, rank)
              for t, rank, is_max in zip(
                  tensors, tensor_ranks, is_max_elem
              )
          ]
          return tf.nest.pack_sequence_as(weights, mask)
  
      raise NotImplementedError(
          'Unrecognized or unimplemented "norm_type": %s' % norm_type
      )
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import neural_structured_learning.configs
  import six
  import tensorflow
label_desc: Use image node_id 6 for calling a global function with example usage: maximize_within_unit_norm(weights, norm_type, epsilon) and returns: tf, weights, tf, tf