def random_in_norm_ball()

in neural_structured_learning/lib/utils.py [0:0]


def random_in_norm_ball(tensor_structure, radius, norm_type):
  """Generates a random sample in a norm ball, conforming the given structure.

  This function expects `tensor_structure` to represent a batch of examples,
  where the tensors in the structure are features. The norms are calculated in
  the whole feature space (including all tensors).

  For L-inf norm, each feature dimension can be sampled independently from a
  uniform distribution. For L2 and L1 norms, we use the result (Corollary 4 in
  Barthe et al., 2005) that the first `n` coordinates of a uniformly random
  point on the `p`-norm unit sphere in `R^(n+p)` are uniformly distributed in
  the `p`-norm unit ball in `R^n`. For L2 norm, we draw random points on the
  unit sphere from the standard normal distribution which density is
  proportional to `exp(-x^2)`(Corollary 1 in Harman and Lacko, 2010). For L1
  norm, we draw random points on the 1-norm unit sphere from the standard
  Laplace distribution which density is proportional to `exp(-|x|)`. More
  details on the papers: https://arxiv.org/abs/math/0503650 and
  https://doi.org/10.1016/j.jmva.2010.06.002

  Args:
    tensor_structure:  A (nested) collection of batched, floating-point tensors.
      Only the structure (number of tensors, shape, and dtype) is used, not the
      values. The first dimension of each tensor is the batch size and must all
      be equal.
    radius: The radius of the norm ball to sample from.
    norm_type: One of `nsl.configs.NormType` which specifies the norm of the
      norm ball.

  Returns:
    A (nested) collection of tensors in the same structure as `tensor_structure`
    but has its values drawn randomly. The norm of each example in the batch
    is less than or equal to the given radius.
  """
  tensors = tf.nest.flatten(tensor_structure)
  if not all(t.dtype.is_floating for t in tensors):
    raise ValueError('random_in_norm_ball() only supports floating tensors.')
  # Reject unless the set of first dimensions is {N}, {None}, or {N, None} for
  # some number N.
  if len(set(t.shape.as_list()[0] for t in tensors) - set([None])) > 1:
    raise ValueError('The first dimension of all tensors should be the same.')
  if isinstance(norm_type, six.string_types):  # Convert string to NormType.
    norm_type = configs.NormType(norm_type)

  if norm_type == configs.NormType.INFINITY:
    return tf.nest.map_structure(
        lambda t: tf.random.uniform(t.shape, -radius, radius, t.dtype),
        tensor_structure)
  elif norm_type == configs.NormType.L2:
    # Sample each coordinate with Normal(0, 1).
    samples = [tf.random.normal(t.shape, dtype=t.dtype) for t in tensors]
    squared_norm = _reduce_across_tensors(tf.reduce_sum,
                                          [tf.square(t) for t in samples])
    # Squared L2 norm of two extra coordinates:
    # Normal(0, 1)^2 + Normal(0, 1)^2 ~ Chi^2(2) ~ Exponential(0.5)
    # Exponential(lambda) ~ Gamma(1, beta=lambda)
    extra = tf.random.gamma(squared_norm.shape, 1.0, 0.5, squared_norm.dtype)
    scale = radius / tf.math.sqrt(squared_norm + extra)
    return tf.nest.pack_sequence_as(
        tensor_structure,
        [t * _expand_to_rank(scale, t.shape.rank) for t in samples])
  elif norm_type == configs.NormType.L1:
    # Sample each coordinate w/ Laplace(0, 1) ~ Exponential(1) - Exponential(1)
    samples = []
    for t in tensors:
      samples.append(
          tf.random.gamma(t.shape, 1., 1., t.dtype) -
          tf.random.gamma(t.shape, 1., 1., t.dtype))
    l1_norm = _reduce_across_tensors(tf.reduce_sum,
                                     [tf.abs(t) for t in samples])
    # L1 norm of one extra coordinate: |Laplace(0, 1)| ~ Exponential(1).
    extra = tf.random.gamma(l1_norm.shape, 1.0, 1.0, l1_norm.dtype)
    scale = radius / (l1_norm + extra)
    return tf.nest.pack_sequence_as(
        tensor_structure,
        [t * _expand_to_rank(scale, t.shape.rank) for t in samples])
  else:
    raise ValueError(f'NormType not supported: {norm_type}')