def make_jd_coroutine(axis_name): def model_coroutine(): w = yield tfd.JointDistributionCoroutine.Root(tfd.Normal(0., 1.)) x = yield sharded.Sharded(tfd.Sample(tfd.Normal(w, 1.), 1), shard_axis_name=axis_name) yield sharded.Sharded(tfd.Independent(tfd.Normal(x, 1.), 1), shard_axis_name=axis_name) return jd.JointDistributionCoroutine(model_coroutine)
def make_jd_coroutine(axis_name): def model_coroutine(): w = yield tfd.JointDistributionCoroutine.Root(tfd.Normal(0., 1.)) x = yield sharded.ShardedSample( tfd.Normal(w, 1.), test_lib.NUM_DEVICES, shard_axis_name=axis_name) yield sharded.ShardedIndependent( tfd.Normal(x, 1.), 1, shard_axis_name=axis_name) return jd.JointDistributionCoroutine( model_coroutine, shard_axis_name=axis_name)
def __init__(self, *args, name='JointDensityCoroutine', **kwargs): """Construct the `JointDensityCoroutine` density. See the documentation for JointDistributionCoroutine Args: *args: Positional arguments forwarded to JointDistributionCoroutine. name: The name for ops managed by the density. Default value: `JointDensityCoroutine`. **kwargs: Named arguments forwarded to JointDistributionCoroutine. """ with tf.name_scope(name) as name: self._joint_distribution_coroutine = jdc_lib.JointDistributionCoroutine( name=name, *args, **kwargs)
NUM_DEVICES = 4 def per_replica_to_tensor(value): return tf.nest.map_structure( lambda per_replica: tf.stack(per_replica.values, axis=0), value) def model_coroutine(): w = yield tfd.JointDistributionCoroutine.Root(tfd.Normal(0., 1.)) x = yield sharded.ShardedSample(tfd.Normal(w, 1.), NUM_DEVICES) yield sharded.ShardedIndependent(tfd.Normal(x, 1.), 1) 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(