tfep.nn.transformers.spline.neural_spline_transformer

tfep.nn.transformers.spline.neural_spline_transformer(x, x0, y0, widths, heights, slopes)[source]

Implement the neural spline transformer.

This is an implementation of the neural spline transformer proposed in [1]. Using the therminology in [1], the spline function is defined from K+1 knots (x, y) that delimit K bins. The domain of the spline for feature i is in the interval x0[i] <= x[i] <= x0[i] + cumsum(widths). All values outside this interval are transformed linearly following the slopes of the first and last knot.

Note that this function implements only the neural spline function. Compared, to NeuralSplineTransformer, it has several differences: - It takes as parameters directly widths, heights, and slopes of the

K+1 knots used for interpolation rather than parameters that are then passed to softmax and softplus functions.

  • Periodic features are not shifted nor mapped to their domain. These are expected to be passed in their domain or they will be transformed linearly.

  • It does not guarantee a minimum bin size and slope for numerical stability.

Parameters:

  • x (torch.Tensor) – Shape (batch_size, n_features). Input features. Periodic features must be within their domain. Non-periodic features outside their domain will be transformed linearly.

  • x0 (torch.Tensor) – Shape (n_features,) or (batch_size, n_features). Position of the first of the K+1 knots determining the positions of the K bins for the input.

  • y0 (torch.Tensor) – Shape (n_features,) or (batch_size, n_features). Position of the first of the K+1 knots determining the positions of the K bins for the output.

  • widths (torch.Tensor) – Shape (batch_size, K, n_features). widths[b, k, i] is the width of the k-th bin between the k-th and (k+1)-th knot for the i-th feature and b-th batch.

  • heights (torch.Tensor) – Shape (batch_size, K, n_features). heights[b, k, i] is the height of the k-th bin between the k-th and (k+1)-th knot for the i-th feature and b-th batch.

  • slopes (torch.Tensor) – Shape (batch_size, K+1, n_features). slopes[b, k, i] is the slope at the k-th knot.

References

[1] Durkan C, Bekasov A, Murray I, Papamakarios G. Neural spline flows.

arXiv preprint arXiv:1906.04032. 2019 Dec 12.