tfp.substrates.jax.bijectors.ScaleMatvecLU

Matrix-vector multiply using LU decomposition.

Inherits From: Bijector

View aliases

Main aliases

tfp.experimental.substrates.jax.bijectors.ScaleMatvecLU

tfp.substrates.jax.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.stateless_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.stateless_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

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., 'ScaleMatvecLU').

Raises
`ValueError`	If both/neither `channels` and `lower_upper`/`permutation` are specified.

Attributes
`dtype`
`forward_min_event_ndims`	Returns the minimal number of dimensions bijector.forward operates on. Multipart bijectors return structured `ndims`, which indicates the expected structure of their inputs. Some multipart bijectors, notably Composites, may return structures of `None`.
`graph_parents`	Returns this `Bijector`'s graph_parents as a Python list.
`has_static_min_event_ndims`	Returns True if the bijector has statically-known `min_event_ndims`.
`inverse_min_event_ndims`	Returns the minimal number of dimensions bijector.inverse operates on. Multipart bijectors return structured `event_ndims`, which indicates the expected structure of their outputs. Some multipart bijectors, notably Composites, may return structures of `None`.
`is_constant_jacobian`	Returns true iff the Jacobian matrix is not a function of x. Note: Jacobian matrix is either constant for both forward and inverse or neither.
`lower_upper`
`name`	Returns the string name of this `Bijector`.
`parameters`	Dictionary of parameters used to instantiate this `Bijector`.
`permutation`
`trainable_variables`
`validate_args`	Returns True if Tensor arguments will be validated.
`variables`

Methods

`forward`

View source

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

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

Args
`x`	`Tensor` (structure). The input to the 'forward' evaluation.
`name`	The name to give this op.
`**kwargs`	Named arguments forwarded to subclass implementation.

Returns
`Tensor` (structure).

Raises
`TypeError`	if `self.dtype` is specified and `x.dtype` is not `self.dtype`.
`NotImplementedError`	if `_forward` is not implemented.

`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.

Args
`input_shape`	`TensorShape` (structure) indicating event-portion shape passed into `forward` function.

Returns
`forward_event_shape_tensor`	`TensorShape` (structure) indicating event-portion shape after applying `forward`. Possibly unknown.

`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.

Args
`input_shape`	`Tensor`, `int32` vector (structure) indicating event-portion shape passed into `forward` function.
`name`	name to give to the op

Returns
`forward_event_shape_tensor`	`Tensor`, `int32` vector (structure) indicating event-portion shape after applying `forward`.

`forward_log_det_jacobian`

View source

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

Returns both the forward_log_det_jacobian.

Args
`x`	`Tensor` (structure). 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. Multipart bijectors require structured event_ndims, such that `rank(y[i]) - rank(event_ndims[i])` is the same for all elements `i` of the structured input. Furthermore, the first `event_ndims[i]` of each `x[i].shape` must be the same for all `i` (broadcasting is not allowed).
`name`	The name to give this op.
`**kwargs`	Named arguments forwarded to subclass implementation.

Returns
`Tensor` (structure), if this bijector is injective. If not injective this is not implemented.

Raises
`TypeError`	if `y`'s dtype is incompatible with the expected output dtype.
`NotImplementedError`	if neither `_forward_log_det_jacobian` nor {`_inverse`, `_inverse_log_det_jacobian`} are implemented, or this is a non-injective bijector.

`inverse`

View source

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

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

Args
`y`	`Tensor` (structure). The input to the 'inverse' evaluation.
`name`	The name to give this op.
`**kwargs`	Named arguments forwarded to subclass implementation.

Returns
`Tensor` (structure), 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 `y`'s structured dtype is incompatible with the expected output dtype.
`NotImplementedError`	if `_inverse` is not implemented.

`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.

Args
`output_shape`	`TensorShape` (structure) indicating event-portion shape passed into `inverse` function.

Returns
`inverse_event_shape_tensor`	`TensorShape` (structure) indicating event-portion shape after applying `inverse`. Possibly unknown.

`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.

Args
`output_shape`	`Tensor`, `int32` vector (structure) indicating event-portion shape passed into `inverse` function.
`name`	name to give to the op

Returns
`inverse_event_shape_tensor`	`Tensor`, `int32` vector (structure) indicating event-portion shape after applying `inverse`.

`inverse_log_det_jacobian`

View source

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` (structure). 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. Multipart bijectors require structured event_ndims, such that `rank(y[i]) - rank(event_ndims[i])` is the same for all elements `i` of the structured input. Furthermore, the first `event_ndims[i]` of each `x[i].shape` must be the same for all `i` (broadcasting is not allowed).
`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 `x`'s dtype is incompatible with the expected inverse-dtype.
`NotImplementedError`	if `_inverse_log_det_jacobian` is not implemented.

`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:

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

Args
`value`	A `tfd.Distribution`, `tfb.Bijector`, or a (structure of) `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 (structure of) `Tensor` computed by `self.forward`.

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.

tfp.substrates.jax.bijectors.ScaleMatvecLU

View aliases

Examples

References

Args

Raises

Attributes

Methods

forward

forward_dtype

forward_event_ndims

forward_event_shape

forward_event_shape_tensor

forward_log_det_jacobian

inverse

inverse_dtype

inverse_event_ndims

inverse_event_shape

inverse_event_shape_tensor

inverse_log_det_jacobian

__call__

Examples

__eq__

`forward`

`forward_dtype`

`forward_event_ndims`

`forward_event_shape`

`forward_event_shape_tensor`

`forward_log_det_jacobian`

`inverse`

`inverse_dtype`

`inverse_event_ndims`

`inverse_event_shape`

`inverse_event_shape_tensor`

`inverse_log_det_jacobian`

`call`

`eq`