Merge pull request #1999 from reubenharry:dirichlet

PiperOrigin-RevId: 752493986

Author: tensorflower-gardener

spinoffs/inference_gym/inference_gym/internal/test_util.py

@@ -223,17 +223,22 @@ def target_log_prob_fn(*x):
     seed = test_util.test_seed(sampler_type='stateless')
   current_state = tf.nest.map_structure(
       lambda b, e: b(  # pylint: disable=g-long-lambda
-          tf.zeros([num_chains] + list(e), dtype=dtype)),
+          tf.zeros([num_chains] + list(b.inverse_event_shape(e)), dtype=dtype)),
+      model.default_event_space_bijector,
+      model.event_shape)
+  step_size = tf.nest.map_structure(
+      lambda b, e: tf.fill(b.inverse_event_shape(e), step_size),
       model.default_event_space_bijector,
       model.event_shape)
 
   # tfp.mcmc only works well with lists.
   current_state = tf.nest.flatten(current_state)
+  step_size = tf.nest.flatten(step_size)
 
   hmc = tfp.mcmc.HamiltonianMonteCarlo(
       target_log_prob_fn=target_log_prob_fn,
       num_leapfrog_steps=num_leapfrog_steps,
-      step_size=[tf.fill(s.shape, step_size) for s in current_state])
+      step_size=step_size)
   hmc = tfp.mcmc.TransformedTransitionKernel(
       hmc, tf.nest.flatten(model.default_event_space_bijector))
   hmc = tfp.mcmc.DualAveragingStepSizeAdaptation(

spinoffs/inference_gym/inference_gym/targets/BUILD

@@ -38,6 +38,7 @@ py_library(
         ":banana",
         ":bayesian_model",
         ":brownian_motion",
+        ":dirichlet",
         ":eight_schools",
         ":ill_conditioned_gaussian",
         ":item_response_theory",
@@ -132,6 +133,29 @@ py_test(
     ],
 )
 
+py_library(
+    name = "dirichlet",
+    srcs = ["dirichlet.py"],
+    deps = [
+        ":model",
+        # tensorflow_probability dep,
+    ],
+)
+
+# py_strict
+py_test(
+    name = "dirichlet_test",
+    srcs = ["dirichlet_test.py"],
+    deps = [
+        ":dirichlet",
+        # absl/testing:parameterized dep,
+        # numpy dep,
+        # tensorflow dep,
+        # tensorflow_probability/python/internal:test_util dep,
+        "//inference_gym/internal:test_util",
+    ],
+)
+
 py_library(
     name = "eight_schools",
     srcs = ["eight_schools.py"],

spinoffs/inference_gym/inference_gym/targets/dirichlet.py

@@ -0,0 +1,106 @@
+# Copyright 2025 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.
+"""Dirichlet model."""
+
+import numpy as np
+import tensorflow.compat.v2 as tf
+import tensorflow_probability as tfp
+from inference_gym.targets import model
+import tensorflow_probability.substrates.numpy as tfp_np
+
+
+tfb = tfp.bijectors
+tfd = tfp.distributions
+
+__all__ = [
+    'Dirichlet',
+]
+
+
+class Dirichlet(model.Model):
+  """Creates a Dirichlet.
+
+  This function produces a Dirichlet distribution. Low concentration parameters
+  (much below 1) produce a distribution that is typically difficult to sample
+  from.
+  """
+
+  def __init__(
+      self,
+      concentration_vector=np.ones(100) * 0.1,
+      dtype=tf.float32,
+      name='dirichlet',
+      pretty_name='Dirichlet',
+  ):
+    """Construct the Dirichlet.
+
+    Args:
+      concentration_vector: The concentration parameters of the Dirichlet
+        distribution.
+      dtype: Dtype to use for floating point quantities.
+      name: Python `str` name prefixed to Ops created by this class.
+      pretty_name: A Python `str`. The pretty name of this model.
+    """
+
+    dirichlet = tfp.distributions.Dirichlet(
+        concentration=tf.cast(concentration_vector, dtype),
+    )
+    dirichlet_np = tfp_np.distributions.Dirichlet(
+        concentration=np.asarray(concentration_vector),
+    )
+
+    sample_transformations = {
+        'identity': model.Model.SampleTransformation(
+            fn=lambda params: params,
+            pretty_name='Identity',
+            ground_truth_mean=dirichlet_np.mean(),
+            ground_truth_standard_deviation=dirichlet_np.stddev(),
+            dtype=dtype,
+        )
+    }
+
+    self._dirichlet = dirichlet
+
+    super(Dirichlet, self).__init__(
+        default_event_space_bijector=tfb.IteratedSigmoidCentered(
+            validate_args=False, name='iterated_sigmoid'
+        ),
+        event_shape=dirichlet.event_shape,
+        dtype=dirichlet.dtype,
+        name=name,
+        pretty_name=pretty_name,
+        sample_transformations=sample_transformations,
+    )
+
+  def _unnormalized_log_prob(self, value):
+    return self._dirichlet.log_prob(value)
+
+  def sample(self, sample_shape=(), seed=None, name='sample'):
+    """Generate samples of the specified shape from the target distribution.
+
+    The returned samples are exact (and independent) samples from the target
+    distribution of this model.
+
+    Note that a call to `sample()` without arguments will generate a single
+    sample.
+
+    Args:
+      sample_shape: 0D or 1D `int32` `Tensor`. Shape of the generated samples.
+      seed: Python integer or `tfp.util.SeedStream` instance, for seeding PRNG.
+      name: Name to give to the prefix the generated ops.
+
+    Returns:
+      samples: a `Tensor` with prepended dimensions `sample_shape`.
+    """
+    return self._dirichlet.sample(sample_shape, seed=seed, name=name)

spinoffs/inference_gym/inference_gym/targets/dirichlet_test.py

@@ -0,0 +1,78 @@
+# Copyright 2025 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.
+# ============================================================================
+from absl.testing import parameterized
+import numpy as np
+import tensorflow.compat.v2 as tf
+
+from tensorflow_probability.python.internal import test_util as tfp_test_util
+from inference_gym.internal import test_util
+from inference_gym.targets import dirichlet
+
+
+@test_util.multi_backend_test(globals(),
+                              'targets.dirichlet_test')
+class DirichletTest(test_util.InferenceGymTestCase):
+
+  @parameterized.parameters(np.float32, np.float64)
+  def testBasic(self, dtype):
+    """Checks that you get finite values given unconstrained samples.
+
+    We check `unnormalized_log_prob` as well as the values of the sample
+    transformations.
+
+    Args:
+      dtype: Dtype to use for floating point computations.
+    """
+    model = dirichlet.Dirichlet(
+        concentration_vector=np.ones(10, dtype=dtype),
+        dtype=dtype
+    )
+    self.validate_log_prob_and_transforms(
+        model,
+        sample_transformation_shapes=dict(identity=[10],),
+        check_ground_truth_mean=True,
+        check_ground_truth_standard_deviation=True,
+        dtype=dtype,
+    )
+
+  def testMC(self):
+    """Checks true samples from the model against the ground truth."""
+    model = dirichlet.Dirichlet(
+        concentration_vector=np.ones(10, dtype=np.float32),
+        dtype=tf.float32,
+    )
+    self.validate_ground_truth_using_monte_carlo(
+        model,
+        num_samples=int(1e6),
+    )
+
+  @test_util.numpy_disable_gradient_test
+  def testHMC(self):
+    """Checks true samples from the model against the ground truth."""
+    model = dirichlet.Dirichlet(
+        concentration_vector=np.ones(10, dtype=np.float32),
+        dtype=np.float32,
+    )
+    self.validate_ground_truth_using_hmc(
+        model,
+        num_chains=4,
+        num_steps=4000,
+        num_leapfrog_steps=10,
+        step_size=0.025,
+    )
+
+
+if __name__ == '__main__':
+  tfp_test_util.main()