Scope for disabling omnistaging?

[machineko]

Are there a possibility for doing smth like =>

with jax.api.disable_jit():
    func(x,y)

but using omnistag scope?

with jax.api.disable_omnistage():
    jit(func)(x,y)

My use case is function dependent on a mask created from not static arguments and jitting it works very well without omnistaging but rest of the code still need to use it :)

[mattjj]

There is such a function, but it's an internal API function, and we might break it in the future. I don't want to cause you more headaches!

As you may have seen in the omnistaging upgrade guide, oftentimes if you import numpy as np and import jax.numpy as jnp, for code you don't want to stage out it can be a better idea to use np rather than jnp. That is, you can control what gets staged out to XLA simply by only using jnp for things that you might want to stage out.

Would that work? Otherwise, can you provide an example?

[machineko]

In my use case it wasn't possible to change the function to use numpy without a high-performance penalty right now i can't find a full example cause I was modifying my code quite extensively past few days :P

But it was based on that =>

def create_mask_from_lens(x: jnp.ndarray, min_len, max_len, dtype: jnp.dtype = jnp.float32):
    max_len = jnp.expand_dims(max_len, -1)
    min_len = jnp.expand_dims(min_len, -1)
    idx = jnp.broadcast_to(jnp.arange(x.shape[1]), (x.shape[0], x.shape[1]))
    return jnp.less(idx, max_len) & jnp.greater_equal(idx, min_len)

After creating a mask from the lens (inside a function) I just pass it thru index_add ops inside a jax loop (or scan dont rember :P) and it was extremely slow without jit and after jitting with omnistaging error's occurs about index :P (jitting without it give extremely good results but I'm using flax so it's not possible to ignore it)

My code right now

from numba import cuda

@cuda.jit
def gpu_add_irfft(arr, tmp_ir, hop_len):
    x, y = cuda.grid(2)
    stridex, stridey = cuda.gridsize(2)
    for i in range(x, tmp_ir.shape[0], stridex):

        for j in range(y, tmp_ir.shape[1], stridey):
            sample = j * hop_len
            index = sample + i
            cuda.atomic.add(arr, index, tmp_ir[i, j])

def inverse_stft_gpu(
        stft_matrix,
        hop_length=256,
        win_length=1024,
        window="hanning",
        dtype=np.float32):
    n_fft = 2 * (stft_matrix.shape[0] - 1)
    n_frames = stft_matrix.shape[1]

    fft_window = get_window(window, win_length)
    fft_window = pad_center(fft_window, n_fft)
    sq_window = fft_window ** 2

    expected_signal_len = n_fft + hop_length * (n_frames - 1)
    tmp_irfft = jnp.expand_dims(fft_window, -1) * jnp.fft.irfft(stft_matrix, axis=0)

    y = cuda.to_device(np.zeros(expected_signal_len, dtype=dtype))
    gpu_add_irfft[(32, 8), (8, 8)](y, tmp_irfft, hop_length)

    win_sum = cuda.to_device(np.zeros(expected_signal_len, dtype=dtype))

    gpu_add_win[(32, 8), (8, 8)](win_sum, sq_window, tmp_irfft.shape[1], hop_length)

    y, win_sum = np.asarray(y), np.asarray(win_sum)
    non_zero = win_sum != 0.0
    y[non_zero] /= win_sum[non_zero]
    return jnp.asarray(y[int(n_fft // 2): -int(n_fft // 2)])

If u got a better idea I'll take it!

[tshadley]

There is such a function, but it's an internal API function, and we might break it in the future. I don't want to cause you more headaches!

In reference to scope for disabling omnistaging, is there still an internal API function and could I ask for details?
The use-case I'm looking for would be solely to skip XLA compiling for quicker startup and debugging in early stages of python code development.