Python tensorflow.compat.v2.reduce_mean() Examples

def get_cross_entropy_loss(learner_agent_output, env_output, actor_agent_output,
                           actor_action, reward_clipping, discounting,
                           baseline_cost, entropy_cost, num_steps):
  """Computes cross entropy loss."""
  del env_output
  del actor_agent_output
  del reward_clipping
  del discounting
  del baseline_cost
  del entropy_cost

  # Align learner output and actor output.
  # NOTE that for a tensor of num_timesteps=3, learner output has output at
  # timesteps [t1, t2, t3] while actor output has output at timesteps [t0, t1,
  # t2]. Hence the need to align before computing cross-entropy loss.
  policy_logits = learner_agent_output.policy_logits[:-1]
  target_actions = actor_action.oracle_next_action_idx[1:]

  cross_entropy = tf.nn.sparse_softmax_cross_entropy_with_logits(
      labels=target_actions, logits=policy_logits)
  tf.summary.scalar(
      'loss/cross_entropy_loss', tf.reduce_mean(cross_entropy), step=num_steps)
  return cross_entropy 

def test_randomized_qmc_basic(self):
    """Tests the randomization of the random.halton sequences."""
    # This test is identical to the example given in Owen (2017), Figure 5.
    dim = 20
    num_results = 2000
    replica = 5
    seed = 121117

    values = []
    for i in range(replica):
      sample, _ = random.halton.sample(dim, num_results=num_results,
                                       seed=seed + i)
      f = tf.reduce_mean(
          input_tensor=tf.reduce_sum(input_tensor=sample, axis=1)**2)
      values.append(self.evaluate(f))
    self.assertAllClose(np.mean(values), 101.6667, atol=np.std(values) * 2) 

def test_multiivariate_xla_compatible(self):
    """Tests that multiivariate GBM sampling is XLA-compatible."""
    corr_matrix = [[1, 0.1], [0.1, 1]]
    process = tff.models.MultivariateGeometricBrownianMotion(
        dim=2, means=0.05, volatilities=[0.1, 0.2], corr_matrix=corr_matrix,
        dtype=tf.float64)
    times = [0.1, 0.5, 1.0]
    initial_state = [1.0, 2.0]
    @tf.function
    def sample_fn():
      return process.sample_paths(
          times=times, initial_state=initial_state, num_samples=10000)
    samples = tf.xla.experimental.compile(sample_fn)[0]
    log_s = tf.math.log(samples)
    mean = tf.reduce_mean(log_s, axis=0)
    expected_mean = ((process._means - process._vols**2 / 2)
                     * np.array(np.expand_dims(times, -1))
                     + np.log(initial_state))
    self.assertAllClose(mean, expected_mean, atol=1e-2, rtol=1e-2) 

def nonneg_crossentropy(expr, target):
  """A cross entropy operator that is appropriate for NQL outputs.

  Query expressions often evaluate to sparse vectors.  This evaluates cross
  entropy safely.

  Args:
    expr: a Tensorflow expression for some predicted values.
    target: a Tensorflow expression for target values.

  Returns:
    Tensorflow expression for cross entropy.
  """
  expr_replacing_0_with_1 = \
     tf.where(expr > 0, expr, tf.ones(tf.shape(input=expr), tf.float32))
  cross_entropies = tf.reduce_sum(
      input_tensor=-target * tf.math.log(expr_replacing_0_with_1), axis=1)
  return tf.reduce_mean(input_tensor=cross_entropies, axis=0) 

def compute_loss(study_loss_types, current_batch_loss_type, agent, agent_state,
                 env_output, actor_agent_output, actor_action, num_steps):
  """Computes loss using given loss type of the batch.

  Args:
    study_loss_types: An iterable of all loss types used in the present study.
    current_batch_loss_type: A tensor of shape [batch_size] with each value a
      valid loss type for the study.
    agent: An instance of `BaseAgent`.
    agent_state: The initial state of the agent.
    env_output: A nested structure of type `EnvOutput`. The tensors are expected
      to have shape [num_timesteps, batch, ...].
    actor_agent_output: A nested structure of type `AgentOutput`. The tensors
      are expected to have shape [num_timesteps, batch, ....]
    actor_action: An instance of `ActorAction` containing indices of the actions
      chosen by actor. The total number of actions available to actor at any
      point is equal to actor_agent_output.policy_logits.shape()[-1].
    num_steps: An int to be used as step arg for summaries.

  Returns:
    A scalar tensor with computed loss.
  """
  learner_agent_output, _ = agent(env_output, agent_state)
  losses_dict = {}

  for loss_type in study_loss_types:
    losses_dict[loss_type] = LOSS_FNS_REGISTRY[loss_type](
        learner_agent_output, env_output, actor_agent_output, actor_action,
        FLAGS.reward_clipping, FLAGS.discounting, FLAGS.baseline_cost,
        FLAGS.entropy_cost, num_steps)
  # All the losses are time-major tensors, i.e., of shape [num_timesteps - 1,
  # batch_size]. Convert to batch-major and then choose which loss to apply to
  # each row.
  losses_dict = tf.nest.map_structure(rnn.transpose_batch_time, losses_dict)
  loss = tf.reduce_mean(
      utils.gather_from_dict(losses_dict, current_batch_loss_type))
  tf.summary.scalar('loss/loss', loss, step=num_steps)
  return loss 

def loss_fn(params, inputs, targets):
  predicted = params[0] * inputs + params[1]
  loss = tf.reduce_mean(input_tensor=tf.square(predicted - targets))
  return tf_np.asarray(loss) 

def _compute_policy_gradient_loss(logits, actions, advantages, step):
  cross_entropy = tf.nn.sparse_softmax_cross_entropy_with_logits(
      labels=actions, logits=logits)
  advantages = tf.stop_gradient(advantages)
  policy_gradient_loss = cross_entropy * advantages
  tf.summary.scalar(
      'loss/policy_gradient_loss',
      tf.reduce_mean(policy_gradient_loss),
      step=step)
  return policy_gradient_loss 

def _compute_entropy_loss(logits, step):
  policy = tf.nn.softmax(logits)
  log_policy = tf.nn.log_softmax(logits)
  entropy = -tf.reduce_mean(-policy * log_policy, axis=-1)
  tf.summary.scalar('loss/entropy', tf.reduce_mean(entropy), step=step)
  return entropy 

def _compute_baseline_loss(advantages, step):
  # Loss for the baseline, summed over the time dimension. Multiply by 0.5 to
  # match the standard update rule:
  #   d(loss) / d(baseline) = advantage
  baseline_cost = .5 * tf.square(advantages)
  tf.summary.scalar(
      'loss/baseline_cost', tf.reduce_mean(baseline_cost), step=step)
  return baseline_cost 

def call(self, inputs, lengths, training):
    seq_mask = tf.sequence_mask(
        lengths, inputs.shape[1], dtype=tf.dtypes.float32)
    forward_outputs = self.forward_rnn(inputs, training=training)
    reversed_inputs = tf.reverse_sequence(inputs, lengths, seq_axis=1)
    backward_outputs = self.backward_rnn(reversed_inputs, training=training)
    backward_outputs = tf.reverse_sequence(
        backward_outputs, lengths, seq_axis=1)
    outputs = tf.concat([forward_outputs, backward_outputs], axis=-1)
    outputs = outputs * tf.expand_dims(seq_mask, -1)

    if self.reduce_states:
      outputs = tf.reduce_mean(outputs, axis=1)

    return outputs 

def test_variables_receive_gradients(self):
    df = deep_factorized.DeepFactorized()
    with tf.GradientTape() as tape:
      x = tf.random.normal([20])
      loss = -tf.reduce_mean(df.log_prob(x))
    grads = tape.gradient(loss, df.trainable_variables)
    self.assertLen(grads, 8)
    self.assertNotIn(None, grads) 

def test_variables_receive_gradients(self):
    loc = tf.Variable(tf.ones([2], dtype=tf.float32))
    log_scale = tf.Variable(tf.zeros([2], dtype=tf.float32))
    logit_weight = tf.Variable(tf.constant([.3, .7], dtype=tf.float32))
    with tf.GradientTape() as tape:
      dist = self.dist_cls(
          loc=loc, scale=tf.exp(log_scale), weight=tf.nn.softmax(logit_weight))
      x = tf.random.normal([20])
      loss = -tf.reduce_mean(dist.log_prob(x))
    grads = tape.gradient(loss, [loc, log_scale, logit_weight])
    self.assertLen(grads, 3)
    self.assertNotIn(None, grads) 

def test_variables_receive_gradients(self):
    loc = tf.Variable(1., dtype=tf.float32)
    log_scale = tf.Variable(0., dtype=tf.float32)
    with tf.GradientTape() as tape:
      dist = self.dist_cls(loc=loc, scale=tf.exp(log_scale))
      x = tf.random.normal([20])
      loss = -tf.reduce_mean(dist.log_prob(x))
    grads = tape.gradient(loss, [loc, log_scale])
    self.assertLen(grads, 2)
    self.assertNotIn(None, grads) 

def test_univariate_xla_compatible(self):
    """Tests that univariate GBM sampling is XLA-compatible."""
    process = tff.models.GeometricBrownianMotion(0.05, 0.5, dtype=tf.float64)
    @tf.function
    def sample_fn():
      return process.sample_paths(
          times=[0.1, 0.5, 1.0], initial_state=2.0, num_samples=10000)
    samples = tf.xla.experimental.compile(sample_fn)[0]
    log_s = tf.math.log(samples)
    mean = tf.reduce_mean(log_s, axis=0)
    expected_mean = ((process._mu - process._sigma**2 / 2)
                     * np.array([0.1, 0.5, 1.0]) + np.log(2.))
    self.assertAllClose(tf.squeeze(mean), expected_mean, atol=1e-2, rtol=1e-2) 

def test_multivariate_sample_mean_and_variance(self, dtype):
    """Tests the mean and vol of the univariate GBM sampled paths."""
    means = 0.05
    volatilities = [0.1, 0.2]
    corr_matrix = [[1, 0.1], [0.1, 1]]
    process = tff.models.MultivariateGeometricBrownianMotion(
        dim=2, means=means, volatilities=volatilities, corr_matrix=corr_matrix,
        dtype=dtype)
    times = [0.1, 0.5, 1.0]
    initial_state = [1.0, 2.0]
    samples = process.sample_paths(
        times=times, initial_state=initial_state,
        random_type=tff.math.random.RandomType.SOBOL, num_samples=10000)
    log_s = tf.math.log(samples)
    mean = tf.reduce_mean(log_s, axis=0, keepdims=True)
    var = tf.reduce_mean((log_s - mean)**2, axis=0)
    expected_mean = ((process._means - process._vols**2 / 2)
                     * np.array(np.expand_dims(times, -1))
                     + np.log(initial_state))
    expected_var = process._vols**2 * np.array(np.expand_dims(times, -1))
    with self.subTest("Drift"):
      self.assertAllClose(tf.squeeze(mean), expected_mean, atol=1e-3, rtol=1e-3)
    with self.subTest("Variance"):
      self.assertAllClose(tf.squeeze(var), expected_var, atol=1e-3, rtol=1e-3)
    with self.subTest("Correlations"):
      samples = self.evaluate(samples)
      for i in range(len(times)):
        corr = np.corrcoef(samples[:, i, :], rowvar=False)
      self.assertAllClose(corr, corr_matrix, atol=1e-2, rtol=1e-2) 

def test_normal_integral_mean_and_var_correctly_estimated(self):
    n = int(1000)
    # This test is almost identical to the similarly named test in
    # monte_carlo_test.py. The only difference is that we use the Sobol
    # samples instead of the random samples to evaluate the expectations.
    # MC with pseudo random numbers converges at the rate of 1/ Sqrt(N)
    # (N=number of samples). For QMC in low dimensions, the expected convergence
    # rate is ~ 1/N. Hence we should only need 1e3 samples as compared to the
    # 1e6 samples used in the pseudo-random monte carlo.
    dtype = tf.float64
    mu_p = tf.constant([-1., 1.], dtype=dtype)
    mu_q = tf.constant([0., 0.], dtype=dtype)
    sigma_p = tf.constant([0.5, 0.5], dtype=dtype)
    sigma_q = tf.constant([1., 1.], dtype=dtype)
    p = tfp.distributions.Normal(loc=mu_p, scale=sigma_p)
    q = tfp.distributions.Normal(loc=mu_q, scale=sigma_q)

    cdf_sample = random.sobol.sample(2, n, dtype=dtype)
    q_sample = q.quantile(cdf_sample)

    # Compute E_p[X].
    e_x = tf.reduce_mean(q_sample * p.prob(q_sample) / q.prob(q_sample), 0)

    # Compute E_p[X^2 - E_p[X]^2].
    e_x2 = tf.reduce_mean(q_sample**2 * p.prob(q_sample) / q.prob(q_sample)
                          - e_x**2, 0)
    stddev = tf.sqrt(e_x2)

    # Keep the tolerance levels the same as in monte_carlo_test.py.
    self.assertEqual(p.batch_shape, e_x.shape)
    self.assertAllClose(self.evaluate(p.mean()), self.evaluate(e_x), rtol=0.01)
    self.assertAllClose(
        self.evaluate(p.stddev()), self.evaluate(stddev), rtol=0.02) 

def test_docstring_example(self):
    # Produce the first 1000 members of the Halton sequence in 3 dimensions.
    num_results = 1000
    dim = 3
    sample, params = random.halton.sample(dim, num_results=num_results,
                                          seed=127)

    # Evaluate the integral of x_1 * x_2^2 * x_3^3  over the three dimensional
    # hypercube.
    powers = tf.range(1., limit=dim + 1)
    integral = tf.reduce_mean(
        input_tensor=tf.reduce_prod(input_tensor=sample**powers, axis=-1))
    true_value = 1. / tf.reduce_prod(input_tensor=powers + 1.)

    # Produces a relative absolute error of 1.7%.
    self.assertAllClose(
        self.evaluate(integral), self.evaluate(true_value), rtol=0.02)

    # Now skip the first 1000 samples and recompute the integral with the next
    # thousand samples. The sequence_indices argument can be used to do this.

    sequence_indices = tf.range(
        start=1000, limit=1000 + num_results, dtype=tf.int32)
    sample_leaped, _ = random.halton.sample(
        dim, sequence_indices=sequence_indices, randomization_params=params)

    integral_leaped = tf.reduce_mean(
        input_tensor=tf.reduce_prod(
            input_tensor=sample_leaped**powers, axis=-1))
    self.assertAllClose(
        self.evaluate(integral_leaped), self.evaluate(true_value), rtol=0.05) 

def test_normal_integral_mean_and_var_correctly_estimated(self):
    n = int(1000)
    # This test is almost identical to the similarly named test in
    # monte_carlo_test.py. The only difference is that we use the Halton
    # samples instead of the random samples to evaluate the expectations.
    # MC with pseudo random numbers converges at the rate of 1/ Sqrt(N)
    # (N=number of samples). For QMC in low dimensions, the expected convergence
    # rate is ~ 1/N. Hence we should only need 1e3 samples as compared to the
    # 1e6 samples used in the pseudo-random monte carlo.
    mu_p = tf.constant([-1., 1.], dtype=tf.float64)
    mu_q = tf.constant([0., 0.], dtype=tf.float64)
    sigma_p = tf.constant([0.5, 0.5], dtype=tf.float64)
    sigma_q = tf.constant([1., 1.], dtype=tf.float64)
    p = tfd.Normal(loc=mu_p, scale=sigma_p)
    q = tfd.Normal(loc=mu_q, scale=sigma_q)

    cdf_sample, _ = random.halton.sample(2, num_results=n, dtype=tf.float64,
                                         seed=1729)
    q_sample = q.quantile(cdf_sample)

    # Compute E_p[X].
    e_x = tf.reduce_mean(q_sample * p.prob(q_sample) / q.prob(q_sample), 0)

    # Compute E_p[X^2 - E_p[X]^2].
    e_x2 = tf.reduce_mean(q_sample**2 * p.prob(q_sample) / q.prob(q_sample)
                          - e_x**2, 0)
    stddev = tf.sqrt(e_x2)

    # Keep the tolerance levels the same as in monte_carlo_test.py.
    self.assertEqual(p.batch_shape, e_x.shape)
    self.assertAllClose(self.evaluate(p.mean()), self.evaluate(e_x), rtol=0.01)
    self.assertAllClose(
        self.evaluate(p.stddev()), self.evaluate(stddev), rtol=0.02) 

def testPmean(self):
    if extensions.tpu_devices():
      self.skipTest("pmean for TPU is not supported yet")
    devices = self._get_two_devices(require_same_type=True)

    def reduce_mean(f):
      return extensions.pmean(f)

    data = tf_np.asarray(tf.convert_to_tensor(value=[1, 3]))
    pmapped = extensions.pmap(reduce_mean, devices=devices)
    result = pmapped(data)

    self.assertAllClose(result[0], 2)
    self.assertAllClose(result[1], 2) 

def test_partial_sum_func_qmc(self):
    """Tests the QMC evaluation of (x_j + x_{j+1} ...+x_{n})^2.

    A good test of QMC is provided by the function:

      f(x_1,..x_n, x_{n+1}, ..., x_{n+m}) = (x_{n+1} + ... x_{n+m} - m / 2)^2

    with the coordinates taking values in the unit interval. The mean and
    variance of this function (with the uniform distribution over the
    unit-hypercube) is exactly calculable:

      <f> = m / 12, Var(f) = m (5m - 3) / 360

    The purpose of the "shift" (if n > 0) in the coordinate dependence of the
    function is to provide a test for Halton sequence which exhibit more
    dependence in the higher axes.

    This test confirms that the mean squared error of RQMC estimation falls
    as O(N^(2-e)) for any e>0.
    """
    n, m = 5, 5
    dim = n + m
    num_results_lo, num_results_hi = 500, 5000
    replica = 10
    true_mean = m / 12.
    seed_lo = 1925
    seed_hi = 898128

    def func_estimate(x):
      return tf.reduce_mean(
          input_tensor=tf.math.squared_difference(
              tf.reduce_sum(input_tensor=x[:, -m:], axis=-1), m / 2.))

    estimates = []
    for i in range(replica):
      sample_lo, _ = random.halton.sample(
          dim, num_results=num_results_lo, seed=seed_lo + i)
      sample_hi, _ = random.halton.sample(
          dim, num_results=num_results_hi, seed=seed_hi + i)
      f_lo, f_hi = func_estimate(sample_lo), func_estimate(sample_hi)
      estimates.append((self.evaluate(f_lo), self.evaluate(f_hi)))
    var_lo, var_hi = np.mean((np.array(estimates) - true_mean)**2, axis=0)

    # Expect that the variance scales as N^2 so var_hi / var_lo ~ k / 10^2
    # with k a fudge factor accounting for the residual N dependence
    # of the QMC error and the sampling error.
    log_rel_err = np.log(100 * var_hi / var_lo)
    self.assertAllClose(log_rel_err, 0., atol=1.2) 

def average(a, axis=None, weights=None, returned=False):  # pylint: disable=missing-docstring
  if axis is not None and not isinstance(axis, six.integer_types):
    # TODO(wangpeng): Support tuple of ints as `axis`
    raise ValueError('`axis` must be an integer. Tuple of ints is not '
                     'supported yet. Got type: %s' % type(axis))
  a = array_ops.array(a)
  if weights is None:  # Treat all weights as 1
    if not np.issubdtype(a.dtype, np.inexact):
      a = a.astype(utils.result_type(a.dtype, dtypes.default_float_type()))
    avg = tf.reduce_mean(a.data, axis=axis)
    if returned:
      if axis is None:
        weights_sum = tf.size(a.data)
      else:
        weights_sum = tf.shape(a.data)[axis]
      weights_sum = tf.cast(weights_sum, a.data.dtype)
  else:
    if np.issubdtype(a.dtype, np.inexact):
      out_dtype = utils.result_type(a.dtype, weights)
    else:
      out_dtype = utils.result_type(a.dtype, weights,
                                    dtypes.default_float_type())
    a = array_ops.array(a, out_dtype).data
    weights = array_ops.array(weights, out_dtype).data

    def rank_equal_case():
      tf.debugging.Assert(tf.reduce_all(tf.shape(a) == tf.shape(weights)),
                          [tf.shape(a), tf.shape(weights)])
      weights_sum = tf.reduce_sum(weights, axis=axis)
      avg = tf.reduce_sum(a * weights, axis=axis) / weights_sum
      return avg, weights_sum
    if axis is None:
      avg, weights_sum = rank_equal_case()
    else:
      def rank_not_equal_case():
        tf.debugging.Assert(tf.rank(weights) == 1, [tf.rank(weights)])
        weights_sum = tf.reduce_sum(weights)
        axes = tf.convert_to_tensor([[axis], [0]])
        avg = tf.tensordot(a, weights, axes) / weights_sum
        return avg, weights_sum
      # We condition on rank rather than shape equality, because if we do the
      # latter, when the shapes are partially unknown but the ranks are known
      # and different, utils.cond will run shape checking on the true branch,
      # which will raise a shape-checking error.
      avg, weights_sum = utils.cond(tf.rank(a) == tf.rank(weights),
                                    rank_equal_case, rank_not_equal_case)

  avg = array_ops.array(avg)
  if returned:
    weights_sum = array_ops.broadcast_to(weights_sum, tf.shape(avg.data))
    return avg, weights_sum
  return avg