distributedRandom.py

# contains helper functions for generating random numbers with a given distribution
import tensorflow as tf
import tensorflow_probability as tfp

def uniform(range):
  N = tf.size(range)
  range = tf.reshape(range, (-1,))
  def F(shape):
    idx = tf.random.uniform(shape, minval=0, maxval=N, dtype=tf.int32)
    res = tf.gather(range, idx)
    tf.assert_equal(tf.shape(res), shape)
    return res
  return F

def weighted(weights, values):
  N = tf.size(weights)
  weights = tf.reshape(weights, (1, N))
  weights = tf.cast(weights, tf.float32)
  weights = tf.math.log(weights)
  values = tf.reshape(values, (N,))
  def F(shape):
    idx = tf.random.categorical(weights, tf.reduce_prod(shape))
    res = tf.gather(values, idx)
    res = tf.reshape(res, shape)
    tf.assert_equal(tf.shape(res), shape)
    return res
  return F

def PowerDistribution(power=2.0):
  def F(range):
    N = tf.size(range)
    weights = tf.range(N, dtype=tf.float32)[::-1] + 1
    weights = tf.pow(weights, power)
    return weighted(weights, range)
  return F

def BetaDistribution(alpha=1.0, beta=1.0):
  def F(range):
    N = tf.size(range)
    d = tfp.distributions.Beta(alpha, beta)
    weights = tf.linspace(0.0, 1.0, N)
    weights = d.prob(weights)
    return weighted(weights, range)
  return F

def UPowerDistribution(power=2.0):
  def F(range):
    N = tf.size(range)
    weights = tf.range(N, dtype=tf.float32)[::-1] + 1
    weights = tf.pow(weights, power)
    # combine left and mirrored left
    weights = tf.concat([weights[:N//2], weights[:N//2][::-1]], 0)
    return weighted(weights, range)
  return F

def config_to_distribution(config):
  if 'uniform' == config['name']:
    return uniform
  
  if 'power' == config['name']:
    return PowerDistribution(config['power'])
  
  if 'beta' == config['name']:
    return BetaDistribution(alpha=config['alpha'], beta=config['beta'])
  
  if 'upower' == config['name']:
    return UPowerDistribution(config['power'])
  
  raise ValueError('Invalid distribution: %s' % config['name'])