import tensorflow as tf

  def _scipy_pareto(self, concentration, scale):
    # In scipy pareto is defined with scale = 1, so we need to scale.
    return stats.pareto(concentration, scale=scale)

  def testParetoShape(self):
    scale = tf.constant([2.] * 5)
    concentration = tf.constant([2.] * 5)
    pareto = tfd.Pareto(concentration, scale)

    self.assertEqual(self.evaluate(pareto.batch_shape_tensor()), (5,))
    self.assertEqual(pareto.batch_shape, tf.TensorShape([5]))
    self.assertAllEqual(self.evaluate(pareto.event_shape_tensor()), [])
    self.assertEqual(pareto.event_shape, tf.TensorShape([]))

  def testParetoShapeBroadcast(self):
    scale = tf.constant([[3., 2.]])
    concentration = tf.constant([[4.], [5.], [6.]])
    pareto = tfd.Pareto(concentration, scale)

    self.assertAllEqual(self.evaluate(pareto.batch_shape_tensor()), (3, 2))
    self.assertAllEqual(pareto.batch_shape, tf.TensorShape([3, 2]))
    self.assertAllEqual(self.evaluate(pareto.event_shape_tensor()), [])
    self.assertEqual(pareto.event_shape, tf.TensorShape([]))

