tfp.bijectors.RationalQuadraticSpline

See Stable See Nightly

View source on GitHub

A piecewise rational quadratic spline, as developed in [1].

Inherits From: Bijector

tfp.bijectors.RationalQuadraticSpline(
    bin_widths, bin_heights, knot_slopes, range_min=-1, validate_args=False,
    name=None
)

This transformation represents a monotonically increasing piecewise rational quadratic function. Outside of the bounds of knot_x/knot_y, the transform behaves as an identity function.

Typically this bijector will be used as part of a chain, with splines for trailing x dimensions conditioned on some of the earlier x dimensions, and with the inverse then solved first for unconditioned dimensions, then using conditioning derived from those inverses, and so forth. For example, if we split a 15-D xs vector into 3 components, we may implement a forward and inverse as follows:

nsplits = 3

class SplineParams(tf.Module):

  def __init__(self, nbins=32):
    self._nbins = nbins
    self._built = False
    self._bin_widths = None
    self._bin_heights = None
    self._knot_slopes = None

  def _bin_positions(self, x):
    x = tf.reshape(x, [-1, self._nbins])
    return tf.math.softmax(x, axis=-1) * (2 - self._nbins * 1e-2) + 1e-2

  def _slopes(self, x):
    x = tf.reshape(x, [-1, self._nbins - 1])
    return tf.math.softplus(x) + 1e-2

  def __call__(self, x, nunits):
    if not self._built:
      self._bin_widths = tf.keras.layers.Dense(
          nunits * self._nbins, activation=self._bin_positions, name='w')
      self._bin_heights = tf.keras.layers.Dense(
          nunits * self._nbins, activation=self._bin_positions, name='h')
      self._knot_slopes = tf.keras.layers.Dense(
          nunits * (self._nbins - 1), activation=self._slopes, name='s')
      self._built = True
    return tfb.RationalQuadraticSpline(
        bin_widths=self._bin_widths(x),
        bin_heights=self._bin_heights(x),
        knot_slopes=self._knot_slopes(x))

xs = np.random.randn(1, 15).astype(np.float32)  # Keras won't Dense(.)(vec).
splines = [SplineParams() for _ in range(nsplits)]

def spline_flow():
  stack = tfb.Identity()
  for i in range(nsplits):
    stack = tfb.RealNVP(5 * i, bijector_fn=splines[i])(stack)
  return stack

ys = spline_flow().forward(xs)
ys_inv = spline_flow().inverse(ys)  # ys_inv ~= xs

For a one-at-a-time autoregressive flow as in [1], it would be profitable to implement a mask over xs to parallelize either the inverse or the forward pass and implement the other using a tf.while_loop. See tfp.bijectors.MaskedAutoregressiveFlow for support doing so (paired with tfp.bijectors.Invert depending which direction should be parallel).

References

[1]: Conor Durkan, Artur Bekasov, Iain Murray, George Papamakarios. Neural Spline Flows. arXiv preprint arXiv:1906.04032, 2019. https://arxiv.org/abs/1906.04032

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

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, name='forward_dtype', **kwargs
)

Returns the dtype of the output of the forward transformation.

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, 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, name='inverse_dtype', **kwargs
)

Returns the dtype of the output of the inverse transformation.

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

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.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.])