inconsistent results of shift_operator

[geggo]

jaxdf.operators.functions.shift_operator shows inconsistent results when applied to FouriereSeries or FiniteDifferences fields.

how to reproduce
At this gist
there is a jupyter notebook to demonstrate the behaviour.

Expected behavior
No warnings emitted.

The meaning of the dx argument to shift_operator is unclear to me. My naive expectation is that dx describes the amount of shift, i.e. for fs = shift_operator(f, dx) I would expect $fs(\vec{x}) = f(\vec{x} - \vec{dx})$.

Instead, for a vector field the intent seems to perform shifts if the n-th field component along the n-th axis with distance dx[n]. This makes sense to undo shifts occuring e.g. when calculating gradients on staggered grid.

At least I would expect consistent results for FourierSeries and FiniteDifferences fields.

What is the intended behaviour?

Additional context
I would like to use it for correctly calculating e.g. rho*gradient(rho) on a staggered grid, for acoustics simulations using jwave.