Transform

class tensorplay.distributions.Transform(cache_size: int = 0)[source]

Abstract class for invertible transformations with computable log det jacobians. They are primarily used in tensorplay.distributions.TransformedDistribution.

Caching is useful for transforms whose inverses are either expensive or numerically unstable. Note that care must be taken with memoized values since the autograd graph may be reversed. For example while the following works with or without caching:

y = t(x)
t.log_abs_det_jacobian(x, y).backward()  # x will receive gradients.

However the following will error when caching due to dependency reversal:

y = t(x)
z = t.inv(y)
grad(z.sum(), [y])  # error because z is x

Derived classes should implement one or both of _call() or _inverse(). Derived classes that set bijective=True should also implement log_abs_det_jacobian().

Parameters:

cache_size (int) – Size of cache. If zero, no caching is done. If one, the latest single value is cached. Only 0 and 1 are supported.

Variables:

• domain (Constraint) – The constraint representing valid inputs to this transform.

• codomain (Constraint) – The constraint representing valid outputs to this transform which are inputs to the inverse transform.

• bijective (bool) – Whether this transform is bijective. A transform t is bijective iff t.inv(t(x)) == x and t(t.inv(y)) == y for every x in the domain and y in the codomain. Transforms that are not bijective should at least maintain the weaker pseudoinverse properties t(t.inv(t(x)) == t(x) and t.inv(t(t.inv(y))) == t.inv(y).

• sign (int or Tensor) – For bijective univariate transforms, this should be +1 or -1 depending on whether transform is monotone increasing or decreasing.

forward_shape(shape)[source]

Infers the shape of the forward computation, given the input shape. Defaults to preserving shape.

property inv: Transform

Returns the inverse Transform of this transform. This should satisfy t.inv.inv is t.

inverse_shape(shape)[source]

Infers the shapes of the inverse computation, given the output shape. Defaults to preserving shape.

log_abs_det_jacobian(x, y)[source]

Computes the log det jacobian log |dy/dx| given input and output.

property sign: int

Returns the sign of the determinant of the Jacobian, if applicable. In general this only makes sense for bijective transforms.