`AssertionError` due to negative `discriminant`

[Fabian-Unruh]

I am using a rational quadratic spline flow to calculate the posterior distribution of two random variables conditioned on 220 random variables. The training and the results for most data points work fine, but seldom an error is obtained when sampling from the NF. The error is an AssertionError caused by some elements of the discriminant being negative:

File "/home/XXX/venvs/NuclArchMain-py311/lib/python3.11/site-packages/nflows/distributions/base.py", line 65, in sample
    return self._sample(num_samples, context)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/XXX/venvs/NuclArchMain-py311/lib/python3.11/site-packages/nflows/flows/base.py", line 54, in _sample
    samples, _ = self._transform.inverse(noise, context=embedded_context)
                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/XXX/venvs/NuclArchMain-py311/lib/python3.11/site-packages/nflows/transforms/base.py", line 60, in inverse
    return self._cascade(inputs, funcs, context)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/XXX/venvs/NuclArchMain-py311/lib/python3.11/site-packages/nflows/transforms/base.py", line 50, in _cascade
    outputs, logabsdet = func(outputs, context)
                         ^^^^^^^^^^^^^^^^^^^^^^
  File "/home/XXX/venvs/NuclArchMain-py311/lib/python3.11/site-packages/nflows/transforms/coupling.py", line 119, in inverse
    transform_split, logabsdet_split = self._coupling_transform_inverse(
                                       ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/XXX/venvs/NuclArchMain-py311/lib/python3.11/site-packages/nflows/transforms/coupling.py", line 197, in _coupling_transform_inverse
    return self._coupling_transform(inputs, transform_params, inverse=True)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/XXX/venvs/NuclArchMain-py311/lib/python3.11/site-packages/nflows/transforms/coupling.py", line 211, in _coupling_transform
    outputs, logabsdet = self._piecewise_cdf(inputs, transform_params, inverse)
                         ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/XXX/venvs/NuclArchMain-py311/lib/python3.11/site-packages/nflows/transforms/coupling.py", line 492, in _piecewise_cdf
    return spline_fn(
           ^^^^^^^^^^
  File "/home/XXX/venvs/NuclArchMain-py311/lib/python3.11/site-packages/nflows/transforms/splines/rational_quadratic.py", line 46, in unconstrained_rational_quadratic_spline
    ) = rational_quadratic_spline(
        ^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/XXX/venvs/NuclArchMain-py311/lib/python3.11/site-packages/nflows/transforms/splines/rational_quadratic.py", line 135, in rational_quadratic_spline
    assert (discriminant >= 0).all()
           ^^^^^^^^^^^^^^^^^^^^^^^^^
AssertionError

As far as I know, the discriminant should always be zero or positive here and the error can only occur due to numerics. Are there any approaches to deal with this problem?

(My model is the following:)

Flow(
  (_transform): CompositeTransform(
    (_transforms): ModuleList(
      (0-9): 10 x PiecewiseRationalQuadraticCouplingTransform(
        (transform_net): NN_withContext(
          (model): Sequential(
            (Layer 1): Linear(in_features=221, out_features=256, bias=True)
            (Activation function 1): ReLU()
            (Layer 2): Linear(in_features=256, out_features=128, bias=True)
            (Activation function 2): ReLU()
            (Layer 3): Linear(in_features=128, out_features=29, bias=True)
          )
        )
      )
    )
  )
  (_distribution): StandardNormal()
  (_embedding_net): Identity()
)

[ImahnShekhzadeh]

Hi @Fabian-Unruh,

Yes, you are right that this happens because of numerical imprecisions. I suspect that you are using float32 as dtype, which is the default in PyTorch? When I faced this problem, I ended up switching to float64 precision, which solved the numerical imprecision. For inference, i.e. generation, write

if generate:
    torch.set_default_dtype(torch.float64)

For the training ("forward") direction, float32 worked fine for me.

I hope this helps.