def testConvertToNestedTensorRaises_incompatible_dtype(self):
# Dtype checks are strict if TF backend
with self.assertRaisesRegex(
ValueError, 'Tensor conversion requested dtype float64'):
nest_util.convert_to_nested_tensor(value=tf.constant(
1, tf.float32),
dtype=tf.float64)
def testConvertToNestedTensorRaises(
self, value, dtype=None, dtype_hint=None, error=None):
# Convert specs to tensors
def maybe_spec_to_tensor(x):
if isinstance(x, TensorSpec):
return tf.zeros(x.shape, x.dtype, name='tensor')
return x
value = nest.map_structure(maybe_spec_to_tensor, value)
# Structure must be exact.
with self.assertRaisesRegex(*error):
nest_util.convert_to_nested_tensor(
value=value,
dtype=dtype,
dtype_hint=dtype_hint)
def _value(self, dtype=None, name=None, as_ref=False):
"""Get the value returned by `tf.convert_to_tensor(distribution)`.
Note: this function may mutate the distribution instance state by caching
the concretized `Tensor` value.
Args:
dtype: Must return a `Tensor` with the given `dtype` if specified.
name: If the conversion function creates a new `Tensor`, it should use the
given `name` if specified.
as_ref: `as_ref` is true, the function must return a `Tensor` reference,
such as a `Variable`.
Returns:
concretized_distribution_value: `Tensor` identical to
`tf.convert_to_tensor(distribution)`.
#### Examples
```python
tfd = tfp.distributions
x = tfd.Normal(0.5, 1).set_tensor_conversion(tfd.Distribution.mean)
x._value()
# ==> tf.convert_to_tensor(x) ==> 0.5
x._value() + 2
# ==> tf.convert_to_tensor(x) + 2. ==> 2.5
x + 2
# ==> tf.convert_to_tensor(x) + 2. ==> 2.5
```
"""
if as_ref:
raise NotImplementedError(
'Cannot convert a `Distribution` to a reference '
'(e.g., `tf.Variable`).')
if self._concrete_value is None:
if self._convert_to_tensor_fn is None:
raise NotImplementedError(
'Failed to convert object of type {} to Tensor. Contents: {}. '
'Call `distribution.set_tensor_conversion(lambda self: ...)` to '
'enable `tf.convert_to_tensor` capability. For example: '
'`x = tfd.Normal(0,1).set_tensor_conversion(tfd.Distribution.mean)`'
' results in `tf.convert_to_tensor(x)` being identical to '
'`x.mean()`.'.format(type(self), self))
with self._name_and_control_scope('value'):
self._concrete_value = (self._convert_to_tensor_fn(
self.tensor_distribution)
if callable(self._convert_to_tensor_fn)
else self._convert_to_tensor_fn)
if (not tf.is_tensor(self._concrete_value)
and not isinstance(self._concrete_value,
composite_tensor.CompositeTensor)):
self._concrete_value = nest_util.convert_to_nested_tensor( # pylint: disable=protected-access
self._concrete_value,
name=name or 'concrete_value',
dtype=dtype,
dtype_hint=self.tensor_distribution.dtype)
return self._concrete_value
def _call_inverse_log_det_jacobian(self, y, event_ndims, name, **kwargs):
"""Compute inverse_log_det_jacobian over the composition."""
with self._name_and_control_scope(name):
dtype = self.forward_dtype(**kwargs)
y = nest_util.convert_to_nested_tensor(
y, name='y', dtype_hint=dtype,
dtype=None if bijector.SKIP_DTYPE_CHECKS else dtype,
allow_packing=True)
event_ndims = nest_util.coerce_structure(
self.inverse_min_event_ndims, event_ndims)
return self._inverse_log_det_jacobian(y, event_ndims, **kwargs)
def testConvertToNestedTensor(
self, value, dtype=None, dtype_hint=None, name='c2t', expected=None):
# Convert specs to tensors
def maybe_spec_to_tensor(x):
if isinstance(x, TensorSpec):
return tf.zeros(x.shape, x.dtype, name='tensor')
return x
value = nest.map_structure(maybe_spec_to_tensor, value)
# Grab shape/dtype from convert_to_nested_tensor for comparison.
observed = nest.map_structure(
TensorSpec.from_tensor,
nest_util.convert_to_nested_tensor(value, dtype, dtype_hint, name=name))
self.assertAllEqualNested(observed, expected)
def _call_log_prob(self, value, name, **kwargs):
"""Wrapper around _log_prob."""
value = tf.nest.pack_sequence_as(self.dtype, tf.nest.flatten(value))
value = nest_util.convert_to_nested_tensor(
value, name="value", dtype_hint=self.dtype, allow_packing=True
)
with self._name_and_control_scope(name, value, kwargs):
if hasattr(self, "_log_prob"):
return self._log_prob(value, **kwargs)
if hasattr(self, "_prob"):
return tf.math.log(self._prob(value, **kwargs))
raise NotImplementedError(
"log_prob is not implemented: {}".format(type(self).__name__)
)
def _call_inverse_log_det_jacobian(self, y, event_ndims, name, **kwargs):
"""Compute inverse_log_det_jacobian over the composition."""
with self._name_and_control_scope(name):
dtype = self.forward_dtype(**kwargs)
y = nest_util.convert_to_nested_tensor(
y, name='y', dtype_hint=dtype,
dtype=None if bijector.SKIP_DTYPE_CHECKS else dtype,
allow_packing=True)
if event_ndims is None:
if self._has_static_min_event_ndims:
event_ndims = self.inverse_min_event_ndims
else:
raise ValueError('Composition bijector with non-static '
'`min_event_ndims` does not support '
'`event_ndims=None`. Please pass a value '
'for `event_ndims`.')
event_ndims = nest_util.coerce_structure(
self.inverse_min_event_ndims, event_ndims)
return self._inverse_log_det_jacobian(y, event_ndims, **kwargs)
