def iid_sample_fn(*args, **kwargs):
"""Draws iid samples from `fn`."""
with tf.name_scope('iid_sample_fn'):
seed = kwargs.pop('seed', None)
if samplers.is_stateful_seed(seed):
kwargs = dict(kwargs, seed=SeedStream(seed, salt='iid_sample')())
def pfor_loop_body(_):
with tf.name_scope('iid_sample_fn_stateful_body'):
return sample_fn(*args, **kwargs)
else:
# If a stateless seed arg is passed, split it into `n` different
# stateless seeds, so that we don't just get a bunch of copies of the
# same sample.
if not JAX_MODE:
warnings.warn(
'Saw Tensor seed {}, implying stateless sampling. Autovectorized '
'functions that use stateless sampling may be quite slow because '
'the current implementation falls back to an explicit loop. This '
'will be fixed in the future. For now, you will likely see '
'better performance from stateful sampling, which you can invoke '
'by passing a Python `int` seed.'.format(seed))
seed = samplers.split_seed(seed, n=n, salt='iid_sample_stateless')
def pfor_loop_body(i):
with tf.name_scope('iid_sample_fn_stateless_body'):
return sample_fn(*args, seed=tf.gather(seed, i), **kwargs)
draws = parallel_for.pfor(pfor_loop_body, n)
return tf.nest.map_structure(unflatten, draws, expand_composites=True)
def _flat_sample_distributions(self, sample_shape=(), seed=None, value=None):
# This function additionally depends on:
# self._dist_fn_wrapped
# self._dist_fn_args
# self._always_use_specified_sample_shape
num_dists = len(self._dist_fn_wrapped)
if seed is not None and samplers.is_stateful_seed(seed):
seed_stream = SeedStream(seed, salt='JointDistributionSequential')
else:
seed_stream = None
if seed is not None:
seeds = samplers.split_seed(seed, n=num_dists,
salt='JointDistributionSequential')
else:
seeds = [None] * num_dists
ds = []
xs = [None] * num_dists if value is None else list(value)
if len(xs) != num_dists:
raise ValueError('Number of `xs`s must match number of '
'distributions.')
for i, (dist_fn, args) in enumerate(zip(self._dist_fn_wrapped,
self._dist_fn_args)):
ds.append(dist_fn(*xs[:i])) # Chain rule of probability.
# Ensure reproducibility even when xs are (partially) set.
stateful_seed = None if seed_stream is None else seed_stream()
if xs[i] is None:
# TODO(b/129364796): We should ignore args prefixed with `_`; this
# would mean we more often identify when to use `sample_shape=()`
# rather than `sample_shape=sample_shape`.
try: # TODO(b/147874898): Eliminate the stateful fallback 20 Dec 2020.
xs[i] = ds[-1].sample(
() if args and not self._always_use_specified_sample_shape
else sample_shape, seed=seeds[i])
except TypeError as e:
if ('Expected int for argument' not in str(e) and
TENSOR_SEED_MSG_PREFIX not in str(e)) or stateful_seed is None:
raise
if not getattr(self, '_resolving_names', False): # avoid recursion
self._resolving_names = True
resolved_names = self._flat_resolve_names()
self._resolving_names = False
msg = (
'Falling back to stateful sampling for distribution #{i} '
'(0-based) of type `{dist_cls}` with component name '
'"{component_name}" and `dist.name` "{dist_name}". Please '
'update to use `tf.random.stateless_*` RNGs. This fallback may '
'be removed after 20-Dec-2020. ({exc})')
warnings.warn(msg.format(
i=i,
dist_name=ds[-1].name,
component_name=resolved_names[i],
dist_cls=type(ds[-1]),
exc=str(e)))
xs[i] = ds[-1].sample(
() if args and not self._always_use_specified_sample_shape
else sample_shape, seed=stateful_seed)
else:
# This signature does not allow kwarg names. Applies
# `convert_to_tensor` on the next value.
xs[i] = nest.map_structure_up_to(
ds[-1].dtype, # shallow_tree
lambda x, dtype: tf.convert_to_tensor(x, dtype_hint=dtype), # func
xs[i], # x
ds[-1].dtype) # dtype
# Note: we could also resolve distributions up to the first non-`None` in
# `self._model_flatten(value)`, however we omit this feature for simplicity,
# speed, and because it has not yet been requested.
return ds, xs
def _flat_sample_distributions(self,
sample_shape=(),
seed=None,
value=None):
"""Executes `model`, creating both samples and distributions."""
ds = []
values_out = []
if samplers.is_stateful_seed(seed):
seed_stream = SeedStream(seed, salt='JointDistributionCoroutine')
if not self._stateful_to_stateless:
seed = None
else:
seed_stream = None # We got a stateless seed for seed=.
# TODO(b/166658748): Make _stateful_to_stateless always True (eliminate it).
if self._stateful_to_stateless and (seed is not None or not JAX_MODE):
seed = samplers.sanitize_seed(seed,
salt='JointDistributionCoroutine')
gen = self._model_coroutine()
index = 0
d = next(gen)
if self._require_root and not isinstance(d, self.Root):
raise ValueError('First distribution yielded by coroutine must '
'be wrapped in `Root`.')
try:
while True:
actual_distribution = d.distribution if isinstance(
d, self.Root) else d
ds.append(actual_distribution)
# Ensure reproducibility even when xs are (partially) set. Always split.
stateful_sample_seed = None if seed_stream is None else seed_stream(
)
if seed is None:
stateless_sample_seed = None
else:
stateless_sample_seed, seed = samplers.split_seed(seed)
if (value is not None and len(value) > index
and value[index] is not None):
def convert_tree_to_tensor(x, dtype_hint):
return tf.convert_to_tensor(x, dtype_hint=dtype_hint)
# This signature does not allow kwarg names. Applies
# `convert_to_tensor` on the next value.
next_value = nest.map_structure_up_to(
ds[-1].dtype, # shallow_tree
convert_tree_to_tensor, # func
value[index], # x
ds[-1].dtype) # dtype_hint
else:
try:
next_value = actual_distribution.sample(
sample_shape=sample_shape if isinstance(
d, self.Root) else (),
seed=(stateful_sample_seed
if stateless_sample_seed is None else
stateless_sample_seed))
except TypeError as e:
if ('Expected int for argument' not in str(e)
and TENSOR_SEED_MSG_PREFIX not in str(e)) or (
stateful_sample_seed is None):
raise
msg = (
'Falling back to stateful sampling for distribution #{index} '
'(0-based) of type `{dist_cls}` with component name '
'{component_name} and `dist.name` "{dist_name}". Please '
'update to use `tf.random.stateless_*` RNGs. This fallback may '
'be removed after 20-Dec-2020. ({exc})')
component_name = (joint_distribution_lib.
get_explicit_name_for_component(
ds[-1]))
if component_name is None:
component_name = '[None specified]'
else:
component_name = '"{}"'.format(component_name)
warnings.warn(
msg.format(index=index,
component_name=component_name,
dist_name=ds[-1].name,
dist_cls=type(ds[-1]),
exc=str(e)))
next_value = actual_distribution.sample(
sample_shape=sample_shape if isinstance(
d, self.Root) else (),
seed=stateful_sample_seed)
if self._validate_args:
with tf.control_dependencies(
self._assert_compatible_shape(
index, sample_shape, next_value)):
values_out.append(
tf.nest.map_structure(tf.identity, next_value))
else:
values_out.append(next_value)
index += 1
d = gen.send(next_value)
except StopIteration:
pass
return ds, values_out