def make_jd_named(axis_name): return jd.JointDistributionNamed( # pylint: disable=g-long-lambda dict( w=tfd.Normal(0., 1.), x=lambda w: sharded.Sharded( # pylint: disable=g-long-lambda tfd.Sample(tfd.Normal(w, 1.), 1), shard_axis_name=axis_name), data=lambda x: sharded.Sharded( # pylint: disable=g-long-lambda tfd.Independent(tfd.Normal(x, 1.), 1), shard_axis_name=axis_name), ))
distributions = ( ('coroutine', lambda: jd.JointDistributionCoroutine(model_coroutine)), ( 'sequential', lambda: jd.JointDistributionSequential([ # pylint: disable=g-long-lambda tfd.Normal(0., 1.), lambda w: sharded.ShardedSample(tfd.Normal(w, 1.), NUM_DEVICES), lambda x: sharded.ShardedIndependent(tfd.Normal(x, 1.), 1), ])), ( 'named', lambda: jd.JointDistributionNamed( # pylint: disable=g-long-lambda dict( w=tfd.Normal(0., 1.), x=lambda w: sharded.ShardedSample(tfd.Normal(w, 1.), NUM_DEVICES), data=lambda x: sharded.ShardedIndependent( tfd.Normal(x, 1.), 1), ))), ) @test_util.test_all_tf_execution_regimes class JointDistributionTest(test_util.TestCase): def setUp(self): super(JointDistributionTest, self).setUp() self.strategy = tf.distribute.MirroredStrategy( devices=tf.config.list_logical_devices()) def shard_values(self, values): def value_fn(ctx):