tfp.bijectors.ScaleMatvecLU

View source on

Matrix-vector multiply using LU decomposition.

Inherits From: AutoCompositeTensorBijector, Bijector, AutoCompositeTensor

tfp.bijectors.ScaleMatvecLU(
    lower_upper, permutation, validate_args=False, name=None
)

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 or 'trainable_lu_factorization'):
    event_size = tf.convert_to_tensor(
        event_size, dtype_hint=tf.int32, name='event_size')
    batch_shape = tf.convert_to_tensor(
        batch_shape, dtype_hint=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,
        name='lower_upper')
    # Initialize a non-trainable variable for the permutation indices so
    # that its value isn't re-sampled from run-to-run.
    permutation = tf.Variable(
        initial_value=permutation,
        trainable=False,
        name='permutation')
    return lower_upper, permutation

channels = 3
conv1x1 = tfb.ScaleMatvecLU(*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

a = tf.Module()
b = tf.Module()
c = tf.Module()
a.b = b
b.c = c
list(a.submodules) == [b, c]
True
list(b.submodules) == [c]
True
list(c.submodules) == []
True

Methods

copy

View source

copy(
    **override_parameters_kwargs
)

Creates a copy of the bijector.

experimental_batch_shape

View source

experimental_batch_shape(
    x_event_ndims=None, y_event_ndims=None
)

Returns the batch shape of this bijector for inputs of the given rank.

The batch shape of a bijector decribes the set of distinct transformations it represents on events of a given size. For example: the bijector tfb.Scale([1., 2.]) has batch shape [2] for scalar events (event_ndims = 0), because applying it to a scalar event produces two scalar outputs, the result of two different scaling transformations. The same bijector has batch shape [] for vector events, because applying it to a vector produces (via elementwise multiplication) a single vector output.

Bijectors that operate independently on multiple state parts, such as tfb.JointMap, must broadcast to a coherent batch shape. Some events may not be valid: for example, the bijector tfd.JointMap([tfb.Scale([1., 2.]), tfb.Scale([1., 2., 3.])]) does not produce a valid batch shape when event_ndims = [0, 0], since the batch shapes of the two parts are inconsistent. The same bijector does define valid batch shapes of [], [2], and [3] if event_ndims is [1, 1], [0, 1], or [1, 0], respectively.

Since transforming a single event produces a scalar log-det-Jacobian, the batch shape of a bijector with non-constant Jacobian is expected to equal the shape of forward_log_det_jacobian(x, event_ndims=x_event_ndims) or inverse_log_det_jacobian(y, event_ndims=y_event_ndims), for x or y of the specified ndims.

experimental_batch_shape_tensor

View source

experimental_batch_shape_tensor(
    x_event_ndims=None, y_event_ndims=None
)

Returns the batch shape of this bijector for inputs of the given rank.

The batch shape of a bijector decribes the set of distinct transformations it represents on events of a given size. For example: the bijector tfb.Scale([1., 2.]) has batch shape [2] for scalar events (event_ndims = 0), because applying it to a scalar event produces two scalar outputs, the result of two different scaling transformations. The same bijector has batch shape [] for vector events, because applying it to a vector produces (via elementwise multiplication) a single vector output.

Bijectors that operate independently on multiple state parts, such as tfb.JointMap, must broadcast to a coherent batch shape. Some events may not be valid: for example, the bijector tfd.JointMap([tfb.Scale([1., 2.]), tfb.Scale([1., 2., 3.])]) does not produce a valid batch shape when event_ndims = [0, 0], since the batch shapes of the two parts are inconsistent. The same bijector does define valid batch shapes of [], [2], and [3] if event_ndims is [1, 1], [0, 1], or [1, 0], respectively.

Since transforming a single event produces a scalar log-det-Jacobian, the batch shape of a bijector with non-constant Jacobian is expected to equal the shape of forward_log_det_jacobian(x, event_ndims=x_event_ndims) or inverse_log_det_jacobian(y, event_ndims=y_event_ndims), for x or y of the specified ndims.

experimental_compute_density_correction

View source

experimental_compute_density_correction(
    x, tangent_space, backward_compat=False, **kwargs
)

Density correction for this transformation wrt the tangent space, at x.

Subclasses of Bijector may call the most specific applicable method of TangentSpace, based on whether the transformation is dimension-preserving, coordinate-wise, a projection, or something more general. The backward-compatible assumption is that the transformation is dimension-preserving (goes from R^n to R^n).

forward

View source

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

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

forward_dtype

View source

forward_dtype(
    dtype=UNSPECIFIED, name='forward_dtype', **kwargs
)

Returns the dtype returned by forward for the provided input.

forward_event_ndims

View source

forward_event_ndims(
    event_ndims, **kwargs
)

Returns the number of event dimensions produced by forward.

forward_event_shape

View source

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.

forward_event_shape_tensor

View source

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.

forward_log_det_jacobian

View source

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

Returns both the forward_log_det_jacobian.

inverse

View source

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

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

inverse_dtype

View source

inverse_dtype(
    dtype=UNSPECIFIED, name='inverse_dtype', **kwargs
)

Returns the dtype returned by inverse for the provided input.

inverse_event_ndims

View source

inverse_event_ndims(
    event_ndims, **kwargs
)

Returns the number of event dimensions produced by inverse.

inverse_event_shape

View source

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.

inverse_event_shape_tensor

View source

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.

inverse_log_det_jacobian

View source

inverse_log_det_jacobian(
    y, event_ndims=None, 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).

parameter_properties

View source

@classmethod
parameter_properties(
    dtype=tf.float32
)

Returns a dict mapping constructor arg names to property annotations.

This dict should include an entry for each of the bijector's Tensor-valued constructor arguments.

with_name_scope

@classmethod
with_name_scope(
    method
)

Decorator to automatically enter the module name scope.

class MyModule(tf.Module):
  @tf.Module.with_name_scope
  def __call__(self, x):
    if not hasattr(self, 'w'):
      self.w = tf.Variable(tf.random.normal([x.shape[1], 3]))
    return tf.matmul(x, self.w)

Using the above module would produce tf.Variables and tf.Tensors whose names included the module name:

mod = MyModule()
mod(tf.ones([1, 2]))
<tf.Tensor: shape=(1, 3), dtype=float32, numpy=..., dtype=float32)>
mod.w
<tf.Variable 'my_module/Variable:0' shape=(2, 3) dtype=float32,
numpy=..., dtype=float32)>

__call__

View source

__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)

Examples

sigmoid = tfb.Reciprocal()(
    tfb.Shift(shift=1.)(
      tfb.Exp()(
        tfb.Scale(scale=-1.))))
# ==> `tfb.Chain([
#         tfb.Reciprocal(),
#         tfb.Shift(shift=1.),
#         tfb.Exp(),
#         tfb.Scale(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.])

__eq__

View source

__eq__(
    other
)

Return self==value.

__getitem__

View source

__getitem__(
    slices
)

__iter__

View source

__iter__()