def log_prob(self, given, group_ndims=0, name=None):
with tf.name_scope(name=name,
default_name=get_default_scope_name(
'log_prob', self)):
given = self._distribution._check_input_shape(given)
log_prob = self._distribution._log_prob(given)
return reduce_group_ndims(tf.reduce_sum, log_prob, group_ndims)
def prob(self, given, group_ndims=0, name=None):
with tf.name_scope(name,
default_name=get_default_scope_name('prob', self)):
log_p = self.log_prob(given, group_ndims=group_ndims)
p = tf.exp(log_p)
return FlowDistributionDerivedTensor(tensor=p,
flow_origin=log_p.flow_origin)
def __call__(self, x):
y, _ = self.transform(x=x,
compute_y=True,
compute_log_det=False,
name=get_default_scope_name(
camel_to_underscore(
self.__class__.__name__)))
return y
def transform(self, x, compute_y=True, compute_log_det=True, name=None):
"""
Transform `x` into `y`, and compute the log-determinant of `f` at `x`
(i.e., :math:`\\log \\det \\frac{\\partial f(x)}{\\partial x}`).
Args:
x (Tensor): The samples of `x`.
compute_y (bool): Whether or not to compute :math:`y = f(x)`?
Default :obj:`True`.
compute_log_det (bool): Whether or not to compute the
log-determinant? Default :obj:`True`.
name (str): If specified, will use this name as the TensorFlow
operational name scope.
Returns:
(tf.Tensor, tf.Tensor): `y` and the (maybe summed) log-determinant.
The items in the returned tuple might be :obj:`None`
if corresponding `compute_?` argument is set to :obj:`False`.
Raises:
RuntimeError: If both `compute_y` and `compute_log_det` are set
to :obj:`False`.
"""
if not compute_y and not compute_log_det:
raise ValueError('At least one of `compute_y` and '
'`compute_log_det` should be True.')
x = tf.convert_to_tensor(x)
if not self._has_built:
self.build(x)
x = self._x_input_spec.validate('x', x)
with tf.name_scope(name,
default_name=get_default_scope_name(
'transform', self),
values=[x]):
y, log_det = self._transform(x, compute_y, compute_log_det)
if compute_log_det:
with assert_deps([
assert_log_det_shape_matches_input(
log_det=log_det,
input=x,
value_ndims=self.x_value_ndims)
]) as asserted:
if asserted: # pragma: no cover
log_det = tf.identity(log_det)
if y is not None:
maybe_add_histogram(y, 'y')
y = maybe_check_numerics(y, 'y')
if log_det is not None:
maybe_add_histogram(log_det, 'log_det')
log_det = maybe_check_numerics(log_det, 'log_det')
return y, log_det
def inverse_transform(self,
y,
compute_x=True,
compute_log_det=True,
value_ndims=0,
name=None):
"""
Transform `y` into `x`, and compute the log-determinant of `f^{-1}` at
`y` (i.e.,
:math:`\\log \\det \\frac{\\partial f^{-1}(y)}{\\partial y}`).
Args:
y (Tensor): The samples of `y`.
compute_x (bool): Whether or not to compute :math:`x = f^{-1}(y)`?
Default :obj:`True`.
compute_log_det (bool): Whether or not to compute the
log-determinant? Default :obj:`True`.
value_ndims (int): Number of value dimensions.
`log_det.ndims == y.ndims - value_ndims`.
name (str): If specified, will use this name as the TensorFlow
operational name scope.
Returns:
(tf.Tensor, tf.Tensor): `x` and the (maybe summed) log-determinant.
The items in the returned tuple might be :obj:`None`
if corresponding `compute_?` argument is set to :obj:`False`.
Raises:
RuntimeError: If both `compute_x` and `compute_log_det` are set
to :obj:`False`.
RuntimeError: If the flow is not explicitly invertible.
"""
if not compute_x and not compute_log_det:
raise ValueError('At least one of `compute_x` and '
'`compute_log_det` should be True.')
value_ndims = int(value_ndims)
if value_ndims < 0:
raise ValueError(
'`value_ndims` must be >= 0: got {}'.format(value_ndims))
y = tf.convert_to_tensor(y)
with tf.name_scope(name,
default_name=get_default_scope_name(
'inverse_transform', self),
values=[y]):
x, log_det = self._inverse_transform(
y=y, compute_x=compute_x, compute_log_det=compute_log_det)
if log_det is not None and value_ndims > 0:
log_det = tf.reduce_sum(log_det,
axis=list(range(-value_ndims, 0)))
return x, log_det
def transform(self,
x,
compute_y=True,
compute_log_det=True,
value_ndims=0,
name=None):
"""
Transform `x` into `y`, and compute the log-determinant of `f` at `x`
(i.e., :math:`\\log \\det \\frac{\\partial f(x)}{\\partial x}`).
Args:
x (Tensor): The samples of `x`.
compute_y (bool): Whether or not to compute :math:`y = f(x)`?
Default :obj:`True`.
compute_log_det (bool): Whether or not to compute the
log-determinant? Default :obj:`True`.
value_ndims (int): Number of value dimensions.
`log_det.ndims == x.ndims - value_ndims`.
name (str): If specified, will use this name as the TensorFlow
operational name scope.
Returns:
(tf.Tensor, tf.Tensor): `y` and the (maybe summed) log-determinant.
The items in the returned tuple might be :obj:`None`
if corresponding `compute_?` argument is set to :obj:`False`.
Raises:
RuntimeError: If both `compute_y` and `compute_log_det` are set
to :obj:`False`.
"""
if not compute_y and not compute_log_det:
raise ValueError('At least one of `compute_y` and '
'`compute_log_det` should be True.')
value_ndims = int(value_ndims)
if value_ndims < 0:
raise ValueError(
'`value_ndims` must be >= 0: got {}'.format(value_ndims))
x = tf.convert_to_tensor(x)
with tf.name_scope(name,
default_name=get_default_scope_name(
'transform', self),
values=[x]):
y, log_det = self._transform(x=x,
compute_y=compute_y,
compute_log_det=compute_log_det)
if log_det is not None and value_ndims > 0:
log_det = tf.reduce_sum(log_det,
axis=list(range(-value_ndims, 0)))
return y, log_det
def prob(self, given, group_ndims=0, name=None):
"""
Compute the densities of `x` against the distribution.
Args:
given (Tensor): The samples to be tested.
group_ndims (int or tf.Tensor): If specified, the last `group_ndims`
dimensions of the log-densities will be summed up. (default 0)
name: TensorFlow name scope of the graph nodes.
(default "prob").
Returns:
tf.Tensor: The densities of `given`.
"""
with tf.name_scope(name,
default_name=get_default_scope_name('prob', self)):
return tf.exp(self.log_prob(given, group_ndims=group_ndims))
def local_log_probs(self, names):
"""
Get the log-densities of stochastic nodes.
Args:
names (Iterable[str]): Names of the queried stochastic nodes.
Returns:
list[tf.Tensor]: Log-densities of the queried stochastic nodes.
Raises:
KeyError: If non-exist name is queried.
"""
names = self._check_names_exist(names)
ret = []
for name in names:
ns = '{}.log_prob'.format(get_default_scope_name(name))
ret.append(self._stochastic_tensors[name].log_prob(name=ns))
return ret
def apply(self, input):
"""
Apply the layer on `input`, to produce output.
Args:
input (Tensor or list[Tensor]): The input tensor, or a list of
input tensors.
Returns:
The output tensor, or a list of output tensors.
"""
if is_tensor_object(input) or isinstance(input, np.ndarray):
input = tf.convert_to_tensor(input)
ns_values = [input]
else:
input = [tf.convert_to_tensor(i) for i in input]
ns_values = input
if not self._has_built:
self.build(input)
with tf.name_scope(get_default_scope_name('apply', self),
values=ns_values):
return self._apply(input)
def inverse_transform(self,
y,
compute_x=True,
compute_log_det=True,
name=None):
"""
Transform `y` into `x`, and compute the log-determinant of `f^{-1}` at
`y` (i.e.,
:math:`\\log \\det \\frac{\\partial f^{-1}(y)}{\\partial y}`).
Args:
y (Tensor): The samples of `y`.
compute_x (bool): Whether or not to compute :math:`x = f^{-1}(y)`?
Default :obj:`True`.
compute_log_det (bool): Whether or not to compute the
log-determinant? Default :obj:`True`.
name (str): If specified, will use this name as the TensorFlow
operational name scope.
Returns:
(tf.Tensor, tf.Tensor): `x` and the (maybe summed) log-determinant.
The items in the returned tuple might be :obj:`None`
if corresponding `compute_?` argument is set to :obj:`False`.
Raises:
RuntimeError: If both `compute_x` and `compute_log_det` are set
to :obj:`False`.
RuntimeError: If the flow is not explicitly invertible.
"""
if not self.explicitly_invertible:
raise RuntimeError(
'The flow is not explicitly invertible: {!r}'.format(self))
if not compute_x and not compute_log_det:
raise ValueError('At least one of `compute_x` and '
'`compute_log_det` should be True.')
if not self._has_built:
raise RuntimeError('`inverse_transform` cannot be called before '
'the flow has been built; it can be built by '
'calling `build`, `apply` or `transform`: '
'{!r}'.format(self))
y = tf.convert_to_tensor(y)
y = self._y_input_spec.validate('y', y)
with tf.name_scope(name,
default_name=get_default_scope_name(
'inverse_transform', self),
values=[y]):
x, log_det = self._inverse_transform(y, compute_x, compute_log_det)
if compute_log_det:
with assert_deps([
assert_log_det_shape_matches_input(
log_det=log_det,
input=y,
value_ndims=self.y_value_ndims)
]) as asserted:
if asserted: # pragma: no cover
log_det = tf.identity(log_det)
return x, log_det
def shifted_conv2d(input,
out_channels,
kernel_size,
spatial_shift,
strides=(1, 1),
channels_last=True,
conv_fn=conv2d,
name=None,
scope=None,
**kwargs):
"""
2D convolution with shifted input.
This method first pads `input` according to the `kernel_size` and
`spatial_shift` arguments, then do 2D convolution (using `conv_fn`)
with "VALID" padding.
Args:
input (Tensor): The input tensor, at least 4-d.
out_channels (int): The channel numbers of the output.
kernel_size (int or (int, int)): Kernel size over spatial dimensions.
spatial_shift: The `spatial_shift` should be a tuple with two elements
(corresponding to height and width spatial axes), and the elements
can only be -1, 0 or 1.
If the shift for a specific axis is `-1`, then `kernel_size - 1`
zeros will be padded at the end of that axis.
If the shift is `0`, then `(kernel_size - 1) // 2` zeros will be
padded at the front, and `kernel_size // 2` zeros will be padded
at the end that axis.
Otherwise if the shift is `1`, then `kernel_size + 1` zeros will
be padded at the front of that axis.
strides (int or (int, int)): Strides over spatial dimensions.
channels_last (bool): Whether or not the channel axis is the last
axis in `input`? (i.e., the data format is "NHWC")
conv_fn: The 2D convolution function. (default :func:`conv2d`)
\\**kwargs: Other named parameters passed to `conv_fn`.
Returns:
tf.Tensor: The output tensor.
"""
spatial_shift = tuple(spatial_shift)
if len(spatial_shift) != 2 or \
any(s not in (-1, 0, 1) for s in spatial_shift):
raise TypeError('`spatial_shift` must be a tuple with two elements, '
'and the elements can only be -1, 0 or 1.')
kernel_size = validate_conv2d_size_tuple('kernel_size', kernel_size)
if 'padding' in kwargs:
raise ValueError('`padding` argument is not supported.')
input, _, _ = validate_conv2d_input(input, channels_last=channels_last)
rank = len(get_static_shape(input))
pads = [(0, 0)] * rank
is_shifted_conv2d = False
spatial_start = -3 if channels_last else -2
for i, (ksize, shift) in enumerate(zip(kernel_size, spatial_shift)):
axis = i + spatial_start
if shift == 0:
pads[axis] = ((ksize - 1) // 2, ksize // 2)
elif shift == -1:
pads[axis] = (0, ksize - 1)
is_shifted_conv2d = True
else:
assert (shift == 1)
pads[axis] = (ksize - 1, 0)
is_shifted_conv2d = True
# fast routine: no shift, use ordinary conv_fn with padding == 'SAME'
if not is_shifted_conv2d:
return conv_fn(input=input,
out_channels=out_channels,
kernel_size=kernel_size,
strides=strides,
channels_last=channels_last,
padding='SAME',
scope=scope,
name=name,
**kwargs)
# slow routine: pad and use conv_fn with padding == 'VALID'
with tf.variable_scope(scope, default_name=name or 'shifted_conv2d'):
output = tf.pad(input, pads)
output = conv_fn(input=output,
out_channels=out_channels,
kernel_size=kernel_size,
strides=strides,
channels_last=channels_last,
padding='VALID',
scope=get_default_scope_name(
getattr(conv_fn, '__name__', None) or 'conv_fn'),
**kwargs)
return output