TensorFlow 2.0 Beta is available Learn more

tfp.bijectors.MatvecLU

Class MatvecLU

Matrix-vector multiply using LU decomposition.

Inherits From: Bijector

Defined in python/bijectors/matveclu.py.

This bijector is identical to the "Convolution1x1" used in Glow [(Kingma and Dhariwal, 2018)[1].

Examples

Here's an example of initialization via random weights matrix:

def trainable_lu_factorization(
    event_size, batch_shape=(), seed=None, dtype=tf.float32, name=None):
  with tf.name_scope(name, 'trainable_lu_factorization',
                     [event_size, batch_shape]):
    event_size = tf.convert_to_tensor(
        event_size, preferred_dtype=tf.int32, name='event_size')
    batch_shape = tf.convert_to_tensor(
        batch_shape, preferred_dtype=event_size.dtype, name='batch_shape')
    random_matrix = tf.random_uniform(
        shape=tf.concat([batch_shape, [event_size, event_size]], axis=0),
        dtype=dtype,
        seed=seed)
    random_orthonormal = tf.linalg.qr(random_matrix)[0]
    lower_upper, permutation = tf.linalg.lu(random_orthonormal)
    lower_upper = tf.Variable(
        initial_value=lower_upper,
        trainable=True,
        use_resource=True,
        name='lower_upper')
    return lower_upper, permutation

channels = 3
conv1x1 = tfb.MatvecLU(*trainable_lu_factorization(channels),
                       validate_args=True)
x = tf.random_uniform(shape=[2, 28, 28, channels])
fwd = conv1x1.forward(x)
rev_fwd = conv1x1.inverse(fwd)
# ==> x

To initialize this variable outside of TensorFlow, one can also use SciPy, e.g.,

def lu_factorized_random_orthonormal_matrix(channels, dtype=np.float32):
  random_matrix = np.random.rand(channels, channels).astype(dtype)
  lower_upper = scipy.linalg.qr(random_matrix)[0]
  permutation = scipy.linalg.lu(lower_upper, overwrite_a=True)[0]
  permutation = np.argmax(permutation, axis=-2)
  return lower_upper, permutation

References

[1]: Diederik P. Kingma, Prafulla Dhariwal. Glow: Generative Flow with Invertible 1x1 Convolutions. arXiv preprint arXiv:1807.03039, 2018. https://arxiv.org/abs/1807.03039

__init__

__init__(
    lower_upper,
    permutation,
    validate_args=False,
    name=None
)

Creates the MatvecLU bijector.

Args:

  • lower_upper: The LU factorization as returned by tf.linalg.lu.
  • permutation: The LU factorization permutation as returned by tf.linalg.lu.
  • validate_args: Python bool indicating whether arguments should be checked for correctness. Default value: False.
  • name: Python str name given to ops managed by this object. Default value: None (i.e., "MatvecLU").

Raises:

  • ValueError: If both/neither channels and lower_upper/permutation are specified.

Properties

dtype

dtype of Tensors transformable by this distribution.

forward_min_event_ndims

Returns the minimal number of dimensions bijector.forward operates on.

graph_parents

Returns this Bijector's graph_parents as a Python list.

inverse_min_event_ndims

Returns the minimal number of dimensions bijector.inverse operates on.

is_constant_jacobian

Returns true iff the Jacobian matrix is not a function of x.

Returns:

  • is_constant_jacobian: Python bool.

lower_upper

name

Returns the string name of this Bijector.

permutation

validate_args

Returns True if Tensor arguments will be validated.

Methods

__call__

__call__(
    value,
    name=None,
    **kwargs
)

Applies or composes the Bijector, depending on input type.

This is a convenience function which applies the Bijector instance in three different ways, depending on the input:

  1. If the input is a tfd.Distribution instance, return tfd.TransformedDistribution(distribution=input, bijector=self).
  2. If the input is a tfb.Bijector instance, return tfb.Chain([self, input]).
  3. Otherwise, return self.forward(input)

Args:

  • value: A tfd.Distribution, tfb.Bijector, or a Tensor.
  • name: Python str name given to ops created by this function.
  • **kwargs: Additional keyword arguments passed into the created tfd.TransformedDistribution, tfb.Bijector, or self.forward.

Returns:

  • composition: A tfd.TransformedDistribution if the input was a tfd.Distribution, a tfb.Chain if the input was a tfb.Bijector, or a Tensor computed by self.forward.

Examples

sigmoid = tfb.Reciprocal()(
    tfb.AffineScalar(shift=1.)(
      tfb.Exp()(
        tfb.AffineScalar(scale=-1.))))
# ==> `tfb.Chain([
#         tfb.Reciprocal(),
#         tfb.AffineScalar(shift=1.),
#         tfb.Exp(),
#         tfb.AffineScalar(scale=-1.),
#      ])`  # ie, `tfb.Sigmoid()`

log_normal = tfb.Exp()(tfd.Normal(0, 1))
# ==> `tfd.TransformedDistribution(tfd.Normal(0, 1), tfb.Exp())`

tfb.Exp()([-1., 0., 1.])
# ==> tf.exp([-1., 0., 1.])

forward

forward(
    x,
    name='forward',
    **kwargs
)

Returns the forward Bijector evaluation, i.e., X = g(Y).

Args:

  • x: Tensor. The input to the "forward" evaluation.
  • name: The name to give this op.
  • **kwargs: Named arguments forwarded to subclass implementation.

Returns:

Tensor.

Raises:

  • TypeError: if self.dtype is specified and x.dtype is not self.dtype.
  • NotImplementedError: if _forward is not implemented.

forward_event_shape

forward_event_shape(input_shape)

Shape of a single sample from a single batch as a TensorShape.

Same meaning as forward_event_shape_tensor. May be only partially defined.

Args:

  • input_shape: TensorShape indicating event-portion shape passed into forward function.

Returns:

  • forward_event_shape_tensor: TensorShape indicating event-portion shape after applying forward. Possibly unknown.

forward_event_shape_tensor

forward_event_shape_tensor(
    input_shape,
    name='forward_event_shape_tensor'
)

Shape of a single sample from a single batch as an int32 1D Tensor.

Args:

  • input_shape: Tensor, int32 vector indicating event-portion shape passed into forward function.
  • name: name to give to the op

Returns:

  • forward_event_shape_tensor: Tensor, int32 vector indicating event-portion shape after applying forward.

forward_log_det_jacobian

forward_log_det_jacobian(
    x,
    event_ndims,
    name='forward_log_det_jacobian',
    **kwargs
)

Returns both the forward_log_det_jacobian.

Args:

  • x: Tensor. The input to the "forward" Jacobian determinant evaluation.
  • event_ndims: Number of dimensions in the probabilistic events being transformed. Must be greater than or equal to self.forward_min_event_ndims. The result is summed over the final dimensions to produce a scalar Jacobian determinant for each event, i.e. it has shape rank(x) - event_ndims dimensions.
  • name: The name to give this op.
  • **kwargs: Named arguments forwarded to subclass implementation.

Returns:

Tensor, if this bijector is injective. If not injective this is not implemented.

Raises:

  • TypeError: if self.dtype is specified and y.dtype is not self.dtype.
  • NotImplementedError: if neither _forward_log_det_jacobian nor {_inverse, _inverse_log_det_jacobian} are implemented, or this is a non-injective bijector.

inverse

inverse(
    y,
    name='inverse',
    **kwargs
)

Returns the inverse Bijector evaluation, i.e., X = g^{-1}(Y).

Args:

  • y: Tensor. The input to the "inverse" evaluation.
  • name: The name to give this op.
  • **kwargs: Named arguments forwarded to subclass implementation.

Returns:

Tensor, if this bijector is injective. If not injective, returns the k-tuple containing the unique k points (x1, ..., xk) such that g(xi) = y.

Raises:

  • TypeError: if self.dtype is specified and y.dtype is not self.dtype.
  • NotImplementedError: if _inverse is not implemented.

inverse_event_shape

inverse_event_shape(output_shape)

Shape of a single sample from a single batch as a TensorShape.

Same meaning as inverse_event_shape_tensor. May be only partially defined.

Args:

  • output_shape: TensorShape indicating event-portion shape passed into inverse function.

Returns:

  • inverse_event_shape_tensor: TensorShape indicating event-portion shape after applying inverse. Possibly unknown.

inverse_event_shape_tensor

inverse_event_shape_tensor(
    output_shape,
    name='inverse_event_shape_tensor'
)

Shape of a single sample from a single batch as an int32 1D Tensor.

Args:

  • output_shape: Tensor, int32 vector indicating event-portion shape passed into inverse function.
  • name: name to give to the op

Returns:

  • inverse_event_shape_tensor: Tensor, int32 vector indicating event-portion shape after applying inverse.

inverse_log_det_jacobian

inverse_log_det_jacobian(
    y,
    event_ndims,
    name='inverse_log_det_jacobian',
    **kwargs
)

Returns the (log o det o Jacobian o inverse)(y).

Mathematically, returns: log(det(dX/dY))(Y). (Recall that: X=g^{-1}(Y).)

Note that forward_log_det_jacobian is the negative of this function, evaluated at g^{-1}(y).

Args:

  • y: Tensor. The input to the "inverse" Jacobian determinant evaluation.
  • event_ndims: Number of dimensions in the probabilistic events being transformed. Must be greater than or equal to self.inverse_min_event_ndims. The result is summed over the final dimensions to produce a scalar Jacobian determinant for each event, i.e. it has shape rank(y) - event_ndims dimensions.
  • name: The name to give this op.
  • **kwargs: Named arguments forwarded to subclass implementation.

Returns:

  • ildj: Tensor, if this bijector is injective. If not injective, returns the tuple of local log det Jacobians, log(det(Dg_i^{-1}(y))), where g_i is the restriction of g to the ith partition Di.

Raises:

  • TypeError: if self.dtype is specified and y.dtype is not self.dtype.
  • NotImplementedError: if _inverse_log_det_jacobian is not implemented.