Bypassing dynamic array limitation in JAX

[unalmis]

Due to limitations in JAX with dynamic array shape sizes, to interpolate a Fourier series to $Q$ non-uniform points with a type 2 nufft, one is often forced to pad the array Q to size S and instead interpolate to $|S| \sim |Q|^2$ points. So the cost becomes $O(K^d log(K^d) + S \log^d \epsilon^{-1})$ for the forward and similiar for the reverse autodiff gradient which further calls type 1 nuffts.

Given a boolean mask which denotes which points in the input array S are junk, would it be simple to add an option to finufft to skip the spreading computation to the junk points and just return some junk value at those indices? In generic python this would be as simple as this.

If this feature happens to be trivial to add, but would be considered bloat, I would be interested in adding this feature to a forked repository in the future, so it would still be useful to me to know that modification.

[ahbarnett]

Dear Kaya, Thanks for your explanation. However, I am confident that we would not want this complication to be added as a feature in FINUFFT. My reasoning is that the runtime of a NUFFT is >>10x that of plain DRAM access. So the current way to handle your use case is: 1) NUFFT output in temp array of size Q, then 2) copy Q into the respective elements of S. This should not be significantly (a few %) slower than your "ideal" case of a new FINUFFT interface which writes into specific elements of S. Keeping our interface simple is a necessity for maintenance. The performance things to think of here are mostly ordering of writing to DRAM in 2). If you reordered the Q pts sensibly beforehand, then step 2) could be close to a linear ordering. Because of the way cache works, this would help make 2) so fast that it would not affect the overall time. Read about HPC of memory and cache, eg, sec 3 of: https://people.freebsd.org/~lstewart/articles/cpumemory.pdf [One improvement to the interface we have considered (mostly wendazhou) is FINUFFT to report a reordering of the NU pts that it would like, and the user should then sort the data in this way. But this is a different request.] Best, Alex

…

On Thu, Jan 8, 2026 at 7:31 PM Kaya Unalmis ***@***.***> wrote: Due to limitation in JAX with dynamic array shape sizes, to interpolate a Fourier series to $Q$ non-uniform points with a type 2 nufft, one is often forced to pad the array Q to size S and instead interpolate to $|S| \gg |Q|$ points. So the cost becomes $O(K^d log(K^d) + S \log^d \epsilon^{-1})$ for the forward and similiar for the autodiff gradient which further calls type 1 nuffts. Given a boolean mask which denotes which points in the input array S are junk, would it be simple to add an option to finufft to skip the spreading computation to the junk points and just return some junk value at those indices? In generic python this would be as simple as: def nufft2_only_good_points(good_points_mask, S, *args, **kwargs): Q = S[good_points_mask] output = np.zeros(S.shape) output[good_points_mask] = finufft.nufft2(Q, *args, **kwargs) return output If this feature happens to be trivial to add, but would be considered feature bloat, I might be interested in adding this feature to a forked repository in the future, so it would still be useful to me to know that modification. — Reply to this email directly, view it on GitHub <https://urldefense.com/v3/__https://github.com/flatironinstitute/finufft/discussions/783__;!!DSb-azq1wVFtOg!SFfWSF1kYAo5MVmw8ouVtOAsneGQzx6M2eYNtAGoMqtvaZr0DY8I8RfdwftLC7MTdDt4ObF-7sEszWhwu7nLdcnauQsH5XOy$>, or unsubscribe <https://urldefense.com/v3/__https://github.com/notifications/unsubscribe-auth/ACNZRSSOLHKK5KADYTDCQV34F3ZGPAVCNFSM6AAAAACREBSAW6VHI2DSMVQWIX3LMV43ERDJONRXK43TNFXW4OZZGMZDGOJSG4__;!!DSb-azq1wVFtOg!SFfWSF1kYAo5MVmw8ouVtOAsneGQzx6M2eYNtAGoMqtvaZr0DY8I8RfdwftLC7MTdDt4ObF-7sEszWhwu7nLdcnauX2SDtby$> . You are receiving this because you are subscribed to this thread.Message ID: ***@***.***>

-- *-------------------------------------------------------------------~^`^~._.~' |\ Alex Barnett Center for Computational Mathematics, Flatiron Institute | \ http://users.flatironinstitute.org/~ahb 646-876-5942

[unalmis]

I think there might be a misunderstanding, so I am clarifying in that case.

So the current way to handle your use case is:

1. NUFFT output in temp array of size Q, then
2. copy Q into the respective elements of S.

I agree it would be preferable to call finufft with just Q points. The issue is that it is impossible to select those Q points on the Python/JAX side:

# This line of code is impossible to execute in JAX.
# To instantiate the array Q, JAX must know the number of elements in the boolean
# array good_points_mask which are True prior to compiling the code.
Q = S[good_points_mask]  

My question is if there is a way to execute that line of code on finufft's end, since the C programming language should not have this limitation. So if I had a forked finufft repo, would it be feasible to just paste this code somewhere in the c source code?

#include <stdio.h>
#include <stdlib.h>
#include <stdbool.h>

double nufft2_only_good_points(bool* good_points_mask, double* S, int N) {
    // N is the size of the array S
    int Q_size = 0;
    for (int i = 0; i < N; i++) {
        if (good_points_mask[i]) {
            Q_size++;
        }
    }

    double* Q = (double*)malloc(Q_size * sizeof(double));
    double* nufft_result = (double*)malloc(Q_size * sizeof(double));
    double* output = (double*)calloc(N, sizeof(double));

    // Doing Q = S[good_points_mask]
    int j = 0;
    for (int i = 0; i < N; i++) {
        if (good_points_mask[i]) {
            Q[j] = S[i];
            j++;
        }
    }

    nufft_result = finufft.nufft2(Q)

    // output[good_points_mask] = nufft_result
    j = 0;
    for (int i = 0; i < N; i++) {
        if (good_points_mask[i]) {
            output[i] = nufft_result[j];
            j++;
        }
    }

    free(Q);
    free(nufft_result);
    return output;
}

Doing a Nufft like this with Q points should be faster than doing a nufft with S = Q^2 points, which as I understand is what you confirming here:

the runtime of a NUFFT is >>10x that of plain DRAM access.

[ahbarnett]

Hi Kaya, Well, this is a request for the jax wrapper, not for FINUFFT itself. The fact that one cannot jax compile Q = S[good_points_mask] is a jax problem that seems way more basic than FINUFFT. You should ask @lgarrison. Google suggests modifying selected array elements via: https://docs.jax.dev/en/latest/_autosummary/jax.Array.at.html#jax.Array.at Rewriting FINUFFT source is not the right route for this. Try to think of other ways to do it. Best, Alex

…

On Mon, Jan 12, 2026 at 2:42 PM Kaya Unalmis ***@***.***> wrote: I think there might be a misunderstanding, so I am clarifying in that case. So the current way to handle your use case is: 1. NUFFT output in temp array of size Q, then 2. copy Q into the respective elements of S. I agree it would be preferable to call finufft with just Q points. The issue is that it is impossible to select those Q points on the Python/JAX side: # This line of code is impossible to execute in JAX.# To instantiate the array Q, JAX must know the number of elements in the boolean# array good_points_mask which are True prior to compiling the code.Q = S[good_points_mask] My question is if there is a way to execute that line of code on finufft's end, since the C programming language should not have this limitation. So if I had a forked finufft repo, would it be feasible to just paste this code somewhere in Finufft's source code? #include <stdio.h>#include <stdlib.h> void nufft2_only_good_points(int* good_points_mask, double* S, int N) { int Q_size = 0; for (int i = 0; i < N; i++) { if (good_points_mask[i]) { Q_size++; } } double* Q = (double*)malloc(Q_size * sizeof(double)); double* nufft_result = (double*)malloc(Q_size * sizeof(double)); double* output = (double*)calloc(N, sizeof(double)); // Doing Q = S[good_points_mask] int j = 0; for (int i = 0; i < N; i++) { if (good_points_mask[i]) { Q[j] = S[i]; j++; } } nufft_result = finufft.nufft2(...) // output[good_points_mask] = nufft_result j = 0; for (int i = 0; i < N; i++) { if (good_points_mask[i]) { output[i] = nufft_result[j]; j++; } } } Doing a Nufft like this with Q points should be faster than doing a nufft with S = Q^2 points, which as I understand is what you confirming here: the runtime of a NUFFT is >>10x that of plain DRAM access. — Reply to this email directly, view it on GitHub <https://urldefense.com/v3/__https://github.com/flatironinstitute/finufft/discussions/783*discussioncomment-15478860__;Iw!!DSb-azq1wVFtOg!VRTU6nPy6d6_Tlvl5ejLKHbZYCHSporZt27A_WbFSUtIVGL9_pDi3BbrY3CsbiioatisKTyXupIMVrvLiyAIeOCfUp5X8WE-$>, or unsubscribe <https://urldefense.com/v3/__https://github.com/notifications/unsubscribe-auth/ACNZRSSWZP4N35OEFTJTYZ34GP2J3AVCNFSM6AAAAACREBSAW6VHI2DSMVQWIX3LMV43URDJONRXK43TNFXW4Q3PNVWWK3TUHMYTKNBXHA4DMMA__;!!DSb-azq1wVFtOg!VRTU6nPy6d6_Tlvl5ejLKHbZYCHSporZt27A_WbFSUtIVGL9_pDi3BbrY3CsbiioatisKTyXupIMVrvLiyAIeOCfUl62MfBN$> . You are receiving this because you commented.Message ID: ***@***.*** com>

-- *-------------------------------------------------------------------~^`^~._.~' |\ Alex Barnett Center for Computational Mathematics, Flatiron Institute | \ http://users.flatironinstitute.org/~ahb 646-876-5942

[unalmis]

Thank you for feedback and suggestions Alex.

@lgarrison I think to support this feature we just need to edit the python wrapper in jax-finufft into finufft source code here:

finufft/python/finufft/finufft/_interfaces.py

Line 913 in 214f4b9

def nufft2d2(x,y,f,out=None,eps=1e-6,isign=-1,**kwargs):

In particular,

def nufft2d2(x,y,f,out=None,eps=1e-6,isign=-1,**kwargs):
    points_mask = kwargs.pop("points_mask", None)
    if points_mask is not None:
        out_shape = x.shape
        x = x[points_mask]
        y = y[points_mask]
    
    out = invoke_guru(2,2,x,y,None,out,None,None,None,f,isign,eps,None,**kwargs)
    if points_mask is not None:
        out_big = np.zeros(out_shape)
        out_big[points_mask] = out
        return out_big

    return out

and then another few lines to the type 1 transform for reverse mode differentiation

def nufft2d1(x,y,c,n_modes=None,out=None,eps=1e-6,isign=1,**kwargs):
    points_mask = kwargs.pop("points_mask", None)
    if points_mask is not None:
        out_shape = x.shape
        x = x[points_mask]
        y = y[points_mask]

    out = invoke_guru(2,1,x,y,None,c,None,None,None,out,isign,eps,n_modes,**kwargs)
    if points_mask is not None:
        out_big = np.zeros(out_shape)
        out_big[points_mask] = out
        return out_big

    return out

Would you accept this if I make a pull request in jax-finufft? Otherwise I may fork jax-finufft and finufft, in which case I'd appreciate if you share any potential jax issues that you foresee.

[ahbarnett]

Dear Kaya, I'm not sure why you're trying to change our python wrapper now! The feature you're asking for is really one in jax, and possibly simply the ability to write into arbitrary array elements. We do not want to change FINUFFT in any way to do this. See what Lehman says re jax. Best, Alex

…

On Mon, Jan 12, 2026 at 7:33 PM Kaya Unalmis ***@***.***> wrote: Thank you for feedback and suggestions Alex. @lgarrison <https://urldefense.com/v3/__https://github.com/lgarrison__;!!DSb-azq1wVFtOg!XY_PU3grLaumUDuzASOd2qRJiuFJ3yXS3y4suK8znQ0gHn6S5dfIif4s6989k7uZ6yfjWFAjKbXxBdNbTD-5ZMLq6fsHPvXM$> I think to support this feature we just need to edit the python wrapper in the finufft source code here: https://github.com/flatironinstitute/finufft/blob/214f4b94643f17b649c347b4f185d3a0670acb56/python/finufft/finufft/_interfaces.py#L913 <https://urldefense.com/v3/__https://github.com/flatironinstitute/finufft/blob/214f4b94643f17b649c347b4f185d3a0670acb56/python/finufft/finufft/_interfaces.py*L913__;Iw!!DSb-azq1wVFtOg!XY_PU3grLaumUDuzASOd2qRJiuFJ3yXS3y4suK8znQ0gHn6S5dfIif4s6989k7uZ6yfjWFAjKbXxBdNbTD-5ZMLq6S3hPRy-$> In particular, def nufft2d2(x,y,f,out=None,eps=1e-6,isign=-1,**kwargs): points_mask = kwargs.pop("points_mask", None) if points_mask: out_shape = x.shape x = x[points_mask] y = y[points_mask] out = invoke_guru(2,2,x,y,None,out,None,None,None,f,isign,eps,None,**kwargs) if points_mask: out_big = np.zeros(out_shape) out_big[points_mask] = out return out_big return out and then another 5 lines to the type 1 transform for reverse mode differentiation def nufft2d1(x,y,c,n_modes=None,out=None,eps=1e-6,isign=1,**kwargs): points_mask = kwargs.pop("points_mask", None) if points_mask: out_size = out_shape x = x[points_mask] y = y[points_mask] out = invoke_guru(2,1,x,y,None,c,None,None,None,out,isign,eps,n_modes,**kwargs) if points_mask: out_big = np.zeros(out_shape) out_big[points_mask] = out return out_big return out Would you accept this if I make a pull request in jax-finufft? Otherwise I may fork jax-finufft and finufft, in which case I'd appreciate if you share any potential issues that you foresee. — Reply to this email directly, view it on GitHub <https://urldefense.com/v3/__https://github.com/flatironinstitute/finufft/discussions/783*discussioncomment-15480295__;Iw!!DSb-azq1wVFtOg!XY_PU3grLaumUDuzASOd2qRJiuFJ3yXS3y4suK8znQ0gHn6S5dfIif4s6989k7uZ6yfjWFAjKbXxBdNbTD-5ZMLq6QtgrNfy$>, or unsubscribe <https://urldefense.com/v3/__https://github.com/notifications/unsubscribe-auth/ACNZRSVGVJVNCAJFRUQPBWT4GQ4M3AVCNFSM6AAAAACREBSAW6VHI2DSMVQWIX3LMV43URDJONRXK43TNFXW4Q3PNVWWK3TUHMYTKNBYGAZDSNI__;!!DSb-azq1wVFtOg!XY_PU3grLaumUDuzASOd2qRJiuFJ3yXS3y4suK8znQ0gHn6S5dfIif4s6989k7uZ6yfjWFAjKbXxBdNbTD-5ZMLq6QgKg2iD$> . You are receiving this because you commented.Message ID: ***@***.*** com>

-- *-------------------------------------------------------------------~^`^~._.~' |\ Alex Barnett Center for Computational Mathematics, Flatiron Institute | \ http://users.flatironinstitute.org/~ahb 646-876-5942

[unalmis]

I understand; jax-finufft wraps those python wrappers with their functions. I simply shared that example so that the jax-finufft team has an idea on how to modify their wrappers.

[lgarrison]

Could you open an issue in the jax-finufft repo so we can get some expert JAX eyes on it? I mostly work on the finufft/CUDA side of things in that project, so I'm not sure what the idiomatic solution for this kind of problem in JAX is.

[unalmis]

I resolved this in flatironinstitute/jax-finufft#216.