Add complex constraint and real Fourier transform. #1762

tillahoffmann · 2024-03-15T00:14:14Z

This PR

adds a _Complex constraint,
modifies the _Real constraint to inherit from _Complex and check that values are indeed real,
adds RealFastFourierTransform.

The RealFastFourierTransform interface differs slightly from jax.numpy.fft.rfftn to respect that batch dimensions precede event dimensions. I.e., rather than specifying the axes along which to transform, one specifies the number of dimensions to transform akin to reinterpreted_batch_ndims in the IndependentTransform. The interface does not support specifying the norm parameter because jit-ing isn't happy with string arguments (cf. jax-ml/jax#3045).

The motivation for this PR is to implement fast Gaussian processes for stationary kernels (cf. section 4 in https://arxiv.org/pdf/2301.08836v4.pdf).

Thank you for taking the time to review my recent PRs. Let me know if they create too much noise in your inbox.

tillahoffmann · 2024-03-15T14:33:59Z

numpyro/distributions/constraints.py

@@ -638,6 +639,11 @@ def feasible_like(self, prototype):
        return jax.numpy.zeros_like(prototype)


+class _Real(_Complex):
+    def __call__(self, x):
+        return super().__call__(x) & (jax.numpy.isreal(x))


Surprisingly, adding the & leads to some warnings not being emitted in test/test_distributions.py::test_distribution_constraints, but only when jit-ed.

I think isreal is a jax function, which will generate a tracer under jit.

One option is to switch to np like in PositiveDefinite. It is better to have separate instances for real and complex though. I don't think that current constraints work for complex numbers.

fehiepsi · 2024-03-15T16:25:40Z

numpyro/distributions/transforms.py

+    def __init__(
+        self,
+        shape=None,
+        ndims=1,


maybe using event_dim to match patterns in constraints. Could you elaborate on why we need this attribute?

The motivation is to specify the number of dimensions along which to transform, e.g., ndim=1 gives the one-dimensional rfft, ndim=2 gives the two-dimensional rfft, etc. Maybe event_ndims or transform_ndims for consistency with reinterpreted_batch_ndims in IndependentTransform?

Oh, I didn't notice that you are supporting n-dimensional fft. Thanks for explaning!
Re event_ndims: looks better to me. We used event_dim just to be consistent with initial distribution designs.

Actually, I like transform_lengths and transform_ndims better. It is up to you. :)

I've kept it with transform_ndims and transform_shape for now. I was thinking shape could be good because that's the terminology numpy uses (although they just call the parameter s).

fehiepsi · 2024-03-15T16:28:04Z

numpyro/distributions/transforms.py

+    def log_abs_det_jacobian(
+        self, x: jnp.ndarray, y: jnp.ndarray, intermediates: None = None
+    ) -> jnp.ndarray:
+        return 0.0


Could you return the whole batch shape for this?

numpyro/distributions/transforms.py

fehiepsi · 2024-03-15T18:33:10Z

numpyro/distributions/transforms.py

+    def domain(self) -> constraints.Constraint:
+        return constraints._IndependentConstraint(
+            constraints._Real(), self.transform_ndims
+        )


maybe constraints.independent(constraints.real, ...)

fehiepsi · 2024-03-15T18:33:19Z

numpyro/distributions/transforms.py

+    @property
+    def codomain(self) -> constraints.Constraint:
+        return constraints._IndependentConstraint(
+            constraints._Complex(), self.transform_ndims


similar to above

fehiepsi · 2024-03-15T18:34:34Z

numpyro/distributions/transforms.py

+    def log_abs_det_jacobian(
+        self, x: jnp.ndarray, y: jnp.ndarray, intermediates: None = None
+    ) -> jnp.ndarray:
+        return jnp.ones_like(x, shape=x.shape[: -self.transform_ndims])


typo: zeros_like

I think you need to broadcast x.shape[: -self.transform_ndims] and y.shape[: -self.transform_ndims]

fehiepsi · 2024-03-15T18:36:38Z

numpyro/distributions/constraints.py

+class _Real(_SingletonConstraint):
+    def __call__(self, x):
+        # XXX: consider to relax this condition to [-inf, inf] interval
+        return (x == x) & (x != float("inf")) & (x != float("-inf")) & np.isreal(x)


Could you remove np.isreal here? We don't check for real input across all constraints. If x is a tracer, np.isreal might not work.

…orm`.

fehiepsi · 2024-03-15T22:15:19Z

Awesome work!! thanks, Till!

fehiepsi · 2024-03-15T22:20:44Z

numpyro/distributions/transforms.py

+        shape = jnp.broadcast_shapes(
+            x.shape[: -self.transform_ndims], y.shape[: -self.transform_ndims]
+        )
+        return jnp.ones_like(x, shape=shape)


oops, wouldn't this zeros_like?

Yes, sorry. Fixed.

tillahoffmann · 2024-03-16T15:27:21Z

Thank you for the fast review and merging!

I just realized we probably want to add the option to unpack the complex coefficients with shape (..., n // 2 + 1) for a signal of shape (..., n) to a real tensor with shape (..., n) because the rest of the library is designed for real-valued tensors. This would be equivalent to what we're doing in Stan here.

One option is to add an argument unpack (or some more informative name) to the transform to indicate if real or complex values should be returned. Another is to add a separate unpacking transform. What do you think?

fehiepsi · 2024-03-16T15:55:57Z

Interesting, I feel that it's better to have a separate transform for the unpack version.

tillahoffmann · 2024-03-16T16:22:23Z

That could look like

class RealFastFourierUnpackTransform(Transform):
    """
    :param size: Size of the last dimension of the transform, required because the size
        cannot be inferred from the shape of the coefficients.
    """
    def __init__(self, size):
        self.size = size

    ...

If we added it with an unpack parameter, we could infer the size from the argument x. I agree that having a separate transform seems like a good idea, but it would also be nice if one doesn't have to specify the shape ahead of time, e.g., if one wants to use the same transform object on different size input. Not sure which one is the better compromise.

Edit: Or did you mean having a transform (which includes unpacking) that inherits from the transform implemented in this PR?

* Add `complex` constraint. * Add real fast Fourier transform. * Remove redundant domain and codomain definitions. * Use numpy for `isreal` check. * Return `log_abs_det_jacobian` with correct batch shape and add test. * Update parameter names for `RealFastFourierTransform`. * Remove `isreal` check. * Broadcast shapes in `log_abs_det_jacobian`. * Update construction of domain and codomain for `RealFastFourierTransform`. * Fix incorrect Jacobian.

tillahoffmann commented Mar 15, 2024

View reviewed changes

fehiepsi reviewed Mar 15, 2024

View reviewed changes

tillahoffmann added 5 commits March 15, 2024 14:13

Add complex constraint.

7d71de5

Add real fast Fourier transform.

35c8005

Remove redundant domain and codomain definitions.

d655956

Use numpy for isreal check.

9a95a1b

Return log_abs_det_jacobian with correct batch shape and add test.

92203cc

tillahoffmann force-pushed the rfft branch 2 times, most recently from 7f89919 to 984f02e Compare March 15, 2024 18:20

Update parameter names for RealFastFourierTransform.

f477a21

tillahoffmann force-pushed the rfft branch from 984f02e to f477a21 Compare March 15, 2024 18:27

fehiepsi reviewed Mar 15, 2024

View reviewed changes

tillahoffmann added 3 commits March 15, 2024 15:02

Remove isreal check.

d738dd5

Broadcast shapes in log_abs_det_jacobian.

dc7b57d

Update construction of domain and codomain for `RealFastFourierTransf…

31d6b84

…orm`.

fehiepsi approved these changes Mar 15, 2024

View reviewed changes

fehiepsi reviewed Mar 15, 2024

View reviewed changes

Fix incorrect Jacobian.

b902f45

fehiepsi merged commit 136f7c0 into pyro-ppl:master Mar 16, 2024
4 checks passed

tillahoffmann deleted the rfft branch March 16, 2024 15:20

This was referenced Mar 20, 2024

Add RecursiveLinearTransform for linear state space models. #1766

Merged

Add ZeroSumNormal distribution #1751

Merged

tillahoffmann mentioned this pull request Nov 18, 2024

Grid restrictions onnela-lab/gptools#25

Open

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Add complex constraint and real Fourier transform. #1762

Add complex constraint and real Fourier transform. #1762

tillahoffmann commented Mar 15, 2024 •

edited

Loading

tillahoffmann Mar 15, 2024

fehiepsi Mar 15, 2024

fehiepsi Mar 15, 2024

tillahoffmann Mar 15, 2024

fehiepsi Mar 15, 2024

fehiepsi Mar 15, 2024 •

edited

Loading

tillahoffmann Mar 15, 2024

fehiepsi Mar 15, 2024

fehiepsi Mar 15, 2024

fehiepsi Mar 15, 2024

fehiepsi Mar 15, 2024

fehiepsi Mar 15, 2024

fehiepsi commented Mar 15, 2024

fehiepsi Mar 15, 2024

tillahoffmann Mar 15, 2024

tillahoffmann commented Mar 16, 2024

fehiepsi commented Mar 16, 2024

tillahoffmann commented Mar 16, 2024 •

edited

Loading

Add complex constraint and real Fourier transform. #1762

Add complex constraint and real Fourier transform. #1762

Conversation

tillahoffmann commented Mar 15, 2024 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

fehiepsi Mar 15, 2024 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

fehiepsi commented Mar 15, 2024

Choose a reason for hiding this comment

Choose a reason for hiding this comment

tillahoffmann commented Mar 16, 2024

fehiepsi commented Mar 16, 2024

tillahoffmann commented Mar 16, 2024 • edited Loading

tillahoffmann commented Mar 15, 2024 •

edited

Loading

fehiepsi Mar 15, 2024 •

edited

Loading

tillahoffmann commented Mar 16, 2024 •

edited

Loading