text: # Copyright 2019 The TensorFlow Probability Authors.
  #
  # Licensed under the Apache License, Version 2.0 (the "License");
  # you may not use this file except in compliance with the License.
  # You may obtain a copy of the License at
  #
  #     http://www.apache.org/licenses/LICENSE-2.0
  #
  # Unless required by applicable law or agreed to in writing, software
  # distributed under the License is distributed on an "AS IS" BASIS,
  # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  # See the License for the specific language governing permissions and
  # limitations under the License.
  # ============================================================================
  """Tests for PlackettLuce distribution."""
  
  from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  
  import numpy as np
  import scipy
  import tensorflow.compat.v1 as tf1
  import tensorflow.compat.v2 as tf
  import tensorflow_probability as tfp
  
  from tensorflow_probability.python.internal import tensorshape_util
  from tensorflow_probability.python.internal import test_util as tfp_test_util
  from tensorflow.python.framework import test_util  # pylint: disable=g-direct-tensorflow-import
  
  tfb = tfp.bijectors
  tfd = tfp.distributions
  
  
  def make_plackett_luce(batch_shape,
                         num_elements,
                         dtype=tf.int32,
                         scores_dtype=tf.float32):
    scores = tf.random.uniform(
        list(batch_shape) + [num_elements], 0.1, 10, dtype=scores_dtype)
    return tfd.PlackettLuce(scores, dtype=dtype, validate_args=True)
  
  
  @test_util.run_all_in_graph_and_eager_modes
  class PlackettLuceTest(tfp_test_util.TestCase):
  
    def setUp(self):
      super(PlackettLuceTest, self).setUp()
      self._rng = np.random.RandomState(42)
  
    def assertRaises(self, error_class, msg):
      if tf.executing_eagerly():
        return self.assertRaisesRegexp(error_class, msg)
      return self.assertRaisesOpError(msg)
  
    def testScores(self):
      s = [0.1, 2., 5.]
      dist = tfd.PlackettLuce(scores=s, validate_args=True)
      self.assertAllClose(s, self.evaluate(dist.scores))
      self.assertAllEqual([3], dist.scores.shape)
  
    def testShapes(self):
      for batch_shape in ([], [1], [2, 3, 4]):
        dist = make_plackett_luce(batch_shape, 10)
        self.assertAllEqual(batch_shape,
                            tensorshape_util.as_list(dist.batch_shape))
        self.assertAllEqual(batch_shape, self.evaluate(dist.batch_shape_tensor()))
        self.assertAllEqual([10], tensorshape_util.as_list(dist.event_shape))
        self.assertAllEqual([10], self.evaluate(dist.event_shape_tensor()))
        # event_shape is available as a constant because the shape is
        # known at graph build time.
        self.assertEqual(10, dist.event_shape)
  
      for batch_shape in ([], [1], [2, 3, 4]):
        dist = make_plackett_luce(batch_shape, tf.constant(10, dtype=tf.int32))
        self.assertAllEqual(
            len(batch_shape), tensorshape_util.rank(dist.batch_shape))
        self.assertAllEqual(batch_shape, self.evaluate(dist.batch_shape_tensor()))
        self.assertAllEqual([10], tensorshape_util.as_list(dist.event_shape))
        self.assertEqual(10, self.evaluate(dist.event_shape_tensor()))
  
    def testDtype(self):
      dist = make_plackett_luce([], 5, dtype=tf.int32)
      self.assertEqual(dist.dtype, tf.int32)
      self.assertEqual(dist.dtype, dist.sample(5).dtype)
      self.assertEqual(dist.dtype, dist.mode().dtype)
      dist = make_plackett_luce([], 5, dtype=tf.int64)
      self.assertEqual(dist.dtype, tf.int64)
      self.assertEqual(dist.dtype, dist.sample(5).dtype)
      self.assertEqual(dist.dtype, dist.mode().dtype)
      self.assertEqual(dist.scores.dtype, tf.float32)
      self.assertEqual(dist.scores.dtype, dist.prob(
          np.array([1, 0, 2, 4, 3], dtype=np.int64)).dtype)
      self.assertEqual(dist.scores.dtype, dist.log_prob(
          np.array([1, 0, 2, 4, 3], dtype=np.int64)).dtype)
      dist = make_plackett_luce([], 5, dtype=tf.int64, scores_dtype=tf.float64)
      self.assertEqual(dist.scores.dtype, tf.float64)
      self.assertEqual(dist.scores.dtype, dist.prob(
          np.array([1, 0, 2, 4, 3], dtype=np.int64)).dtype)
      self.assertEqual(dist.scores.dtype, dist.log_prob(
          np.array([1, 0, 2, 4, 3], dtype=np.int64)).dtype)
  
    def testUnknownShape(self):
      scores = tf1.placeholder_with_default(
          input=[[1e-6, 1000.0], [1000.0, 1e-6]], shape=None)
      dist = tfd.PlackettLuce(scores, validate_args=True)
      sample = dist.sample()
      # Batch entry 0, 1 will sample permutations [1, 0], [0, 1].
      sample_value_batch = self.evaluate(sample)
      self.assertAllEqual([[1, 0], [0, 1]], sample_value_batch)
  
    def testPmfUniform(self):
      # Check that distribution with equal scores gives uniform probability.
      scores = self._rng.uniform() * np.ones(shape=(8, 2, 10), dtype=np.float32)
      dist = tfd.PlackettLuce(scores=scores, validate_args=True)
      np_sample = self.evaluate(dist.sample())
      np_prob = self.evaluate(dist.log_prob(np_sample))
      n_fac = scipy.special.factorial(scores.shape[-1]).astype(np.float32)
      expected_prob = self.evaluate(-tf.math.log(n_fac) * tf.ones_like(np_prob))
      self.assertAllClose(np_prob, expected_prob)
  
    def testPmfMode(self):
      # Check that mode has higher probability than any other random sample.
      scores = self._rng.uniform(low=0.1, size=(10)).astype(np.float32)
      dist = tfd.PlackettLuce(scores=scores, validate_args=True)
      np_sample = self.evaluate(dist.sample())
      np_prob = self.evaluate(dist.log_prob(np_sample))
      mode = self.evaluate(dist.mode())
      mode_prob = self.evaluate(dist.log_prob(mode))
      self.assertLessEqual(np_prob, mode_prob)
  
    def testSample(self):
      scores = np.array([[[0.1, 2.3, 5.], [4.2, 0.5, 3.1]]])
      dist = tfd.PlackettLuce(scores)
      n = 100
      k = scores.shape[-1]
      samples = dist.sample(n, seed=tfp_test_util.test_seed())
      self.assertEqual(samples.dtype, tf.int32)
      sample_values = self.evaluate(samples)
      self.assertAllEqual([n, 1, 2, k], sample_values.shape)
      self.assertFalse(np.any(sample_values < 0))
      self.assertFalse(np.any(sample_values > k))
      self.assertTrue(np.all(np.sum(sample_values, axis=-1) == (k-1)*k//2))
  
  
  @test_util.run_all_in_graph_and_eager_modes
  class PlackettLuceFromVariableTest(tfp_test_util.TestCase):
  
    def testAssertionsProbsAfterMutation(self):
      x = tf.Variable([0.25, 0.25, 0.5])
      d = tfd.PlackettLuce(scores=x, validate_args=True)
      with self.assertRaisesOpError('Condition x > 0 did not hold element-wise'):
        self.evaluate([v.initializer for v in d.variables])
        with tf.control_dependencies([x.assign([-0.25, 0.75, 0.5])]):
          self.evaluate(d.scores_parameter())
  
    def testAssertionsScores(self):
      x = tfp.util.TransformedVariable(0., tfb.Identity(), shape=None)
      with self.assertRaisesRegexp(
          ValueError, 'Argument `scores` must have rank at least 1.'):
        d = tfd.PlackettLuce(scores=x, validate_args=True)
        self.evaluate([v.initializer for v in d.variables])
  
  if __name__ == '__main__':
    tf.test.main()
  