Attend the Women in ML Symposium on December 7 Register now

tfp.bijectors.Blockwise

Stay organized with collections Save and categorize content based on your preferences.

Bijector which applies a list of bijectors to blocks of a Tensor.

Inherits From: Composition, AutoCompositeTensorBijector, Bijector, AutoCompositeTensor

More specifically, given [F_0, F_1, ... F_n] which are scalar or vector bijectors this bijector creates a transformation which operates on the vector [x_0, ... x_n] with the transformation [F_0(x_0), F_1(x_1) ..., F_n(x_n)] where x_0, ..., x_n are blocks (partitions) of the vector.

Example Use:

  blockwise = tfb.Blockwise(
      bijectors=[tfb.Exp(), tfb.Sigmoid()], block_sizes=[2, 1]
    )
  y = blockwise.forward(x)

  # Equivalent to:
  x_0, x_1 = tf.split(x, [2, 1], axis=-1)
  y_0 = tfb.Exp().forward(x_0)
  y_1 = tfb.Sigmoid().forward(x_1)
  y = tf.concat([y_0, y_1], axis=-1)

Keyword arguments can be passed to the inner bijectors by utilizing the inner bijector names, e.g.:

  blockwise = tfb.Blockwise([Bijector1(name='b1'), Bijector2(name='b2')])
  y = blockwise.forward(x, b1={'arg': 1}, b2={'arg': 2})

  # Equivalent to:
  x_0, x_1 = tf.split(x, [1, 1], axis=-1)
  y_0 = Bijector1().forward(x_0, arg=1)
  y_1 = Bijector2().forward(x_1, arg=2)
  y = tf.concat([y_0, y_1], axis=-1)

If every element of the bijectors list is a CompositeTensor, the resulting Blockwise bijector is a CompositeTensor as well. If any element of bijectors is not a CompositeTensor, then a non-CompositeTensor _Blockwise instance is created instead. Bijector subclasses that inherit from Blockwise will also inherit from CompositeTensor.

bijectors A non-empty list of bijectors.
block_sizes A 1-D integer Tensor with each element signifying the length of the block of the input vector to pass to the corresponding bijector. The length of block_sizes must be be equal to the length of bijectors. If left as None, a vector of 1's is used.
validate_args Python bool indicating whether arguments should be checked for correctness.
maybe_changes_size Python bool indicating that this bijector might change the event size. If this is known to be false and set appropriately, then this will lead to improved static shape inference when the block sizes are not statically known.
name Python str, name given to ops managed by this object. Default: E.g., Blockwise([Exp(), Softplus()]).name == 'blockwise_of_exp_and_softplus'.

NotImplementedError If there is a bijector with event_ndims > 1.
ValueError If bijectors list is empty.
ValueError If size of block_sizes does not equal to the length of bijectors or is not a vector.

bijectors

block_sizes

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. (deprecated)

inverse_block_sizes

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.

name Returns the string name of this Bijector.
name_scope Returns a tf.name_scope instance for this class.
non_trainable_variables Sequence of non-trainable variables owned by this module and its submodules.
parameters Dictionary of parameters used to instantiate this Bijector.
submodules Sequence of all sub-modules.

Submodules are modules which are properties of this module, or found as properties of modules which are properties of this module (and so on).

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

trainable_variables Sequence of trainable variables owned by this module and its submodules.

validate_args Returns True if Tensor arguments will be validated.
validate_event_size

variables Sequence of variables owned by this module and its submodules.

Methods

copy

View source

Creates a copy of the bijector.

Args
**override_parameters_kwargs String/value dictionary of initialization arguments to override with new values.

Returns
bijector A new instance of type(self) initialized from the union of self.parameters and override_parameters_kwargs, i.e., dict(self.parameters, **override_parameters_kwargs).

experimental_batch_shape

View source

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.

Args
x_event_ndims Optional Python int (structure) number of dimensions in a probabilistic event passed to forward; this must be greater than or equal to self.forward_min_event_ndims. If None, defaults to self.forward_min_event_ndims. Mutually exclusive with y_event_ndims. Default value: None.
y_event_ndims Optional Python int (structure) number of dimensions in a probabilistic event passed to inverse; this must be greater than or equal to self.inverse_min_event_ndims. Mutually exclusive with x_event_ndims. Default value: None.

Returns
batch_shape TensorShape batch shape of this bijector for a value with the given event rank. May be unknown or partially defined.

experimental_batch_shape_tensor

View source

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.

Args
x_event_ndims Optional Python int (structure) number of dimensions in a probabilistic event passed to forward; this must be greater than or equal to self.forward_min_event_ndims. If None, d