tfp.bijectors.FFJORD

View source on GitHub

Implements a continuous normalizing flow X->Y defined via an ODE.

Inherits From: Bijector

tfp.bijectors.FFJORD( state_time_derivative_fn, ode_solve_fn=None, trace_augmentation_fn=trace_jacobian_hutchinson, initial_time=0.0, final_time=1.0, validate_args=False, dtype=tf.float32, name='ffjord' ) 

This bijector implements a continuous dynamics transformation parameterized by a differential equation, where initial and terminal conditions correspond to domain (X) and image (Y) i.e.

d/dt[state(t)]=state_time_derivative_fn(t, state(t)) state(initial_time) = X state(final_time) = Y 

For this transformation the value of log_det_jacobian follows another differential equation, reducing it to computation of the trace of the jacbian along the trajectory

state_time_derivative = state_time_derivative_fn(t, state(t)) d/dt[log_det_jac(t)] = Tr(jacobian(state_time_derivative, state(t))) 

FFJORD constructor takes two functions ode_solve_fn and trace_augmentation_fn arguments that customize integration of the differential equation and trace estimation.

Differential equation integration is performed by a call to ode_solve_fn. Custom ode_solve_fn must accept the following arguments:

• ode_fn(time, state, **condition_kwargs): Differential equation to be solved. Custom ode_solve_fns may optionally support conditional inputs by accepting a constants dict arg and computing gradients wrt the provided values in **condition_kwargs.
• initial_time: Scalar float or floating Tensor representing the initial time.
• initial_state: Floating Tensor representing the initial state.
• solution_times: 1D floating Tensor of solution times.

And return a Tensor of shape [solution_times.shape, initial_state.shape] representing state values evaluated at solution_times. In addition ode_solve_fn must support nested structures. For more details see the interface of tfp.math.ode.Solver.solve().

Trace estimation is computed simultaneously with state_time_derivative using augmented_state_time_derivative_fn that is generated by trace_augmentation_fn. trace_augmentation_fn takes state_time_derivative_fn, state.shape and state.dtype arguments and returns a augmented_state_time_derivative_fn callable that computes both state_time_derivative and unreduced trace_estimation.

Custom ode_solve_fn and trace_augmentation_fn examples:

# custom_solver_fn: `callable(f, t_initial, t_solutions, y_initial, ...)` # custom_solver_kwargs: Additional arguments to pass to custom_solver_fn. def ode_solve_fn(ode_fn, initial_time, initial_state, solution_times): results = custom_solver_fn(ode_fn, initial_time, solution_times, initial_state, **custom_solver_kwargs) return results ffjord = tfb.FFJORD(state_time_derivative_fn, ode_solve_fn=ode_solve_fn) 

# state_time_derivative_fn: `callable(time, state)` # trace_jac_fn: `callable(time, state)` unreduced jacobian trace function def trace_augmentation_fn(ode_fn, state_shape, state_dtype): def augmented_ode_fn(time, state): return ode_fn(time, state), trace_jac_fn(time, state) return augmented_ode_fn ffjord = tfb.FFJORD(state_time_derivative_fn, trace_augmentation_fn=trace_augmentation_fn) 

For more details on FFJORD and continous normalizing flows see [1], [2].

Usage example:

tfd = tfp.distributions tfb = tfp.bijectors # state_time_derivative_fn: `Callable(time, state)` -> state_time_derivative # e.g. Neural network with inputs and outputs of the same shapes and dtypes. bijector = tfb.FFJORD(state_time_derivative_fn=state_time_derivative_fn) y = bijector.forward(x) # forward mapping x = bijector.inverse(y) # inverse mapping base = tfd.Normal(tf.zeros_like(x), tf.ones_like(x)) # Base distribution transformed_distribution = tfd.TransformedDistribution(base, bijector) 

References

[1]: Chen, T. Q., Rubanova, Y., Bettencourt, J., & Duvenaud, D. K. (2018). Neural ordinary differential equations. In Advances in neural information processing systems (pp. 6571-6583)

[2]: Grathwohl, W., Chen, R. T., Betterncourt, J., Sutskever, I., & Duvenaud, D. (2018). Ffjord: Free-form continuous dynamics for scalable reversible generative models. arXiv preprint arXiv:1810.01367. http://arxiv.org.abs/1810.01367

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