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)