How to optimize performance of large, log-space array updates via JAX?

[Cauch-BS]

Part of my machine-learning code that uses JAX contains some dynamic programming where large array updates occur along the diagonal of a N ✕ N matrix, by use of a lax.scan.

Originally, the original code which I want to 'mimic' was in C++. However, because I wanted the result of this array to be differentiable, and I wanted to take advantage of hardware accelerators, I implemented the original code as a loop in JAX. The original code had $\mathcal{O}(N^3)$ time. I feel that with JAX with sufficient amounts of parallelism, this could be done in $\mathcal{O}(N)$ time.

However, the lax.scan part is somewhat slower than expected –

1. the result is around 10x slower than using C++.
2. using multi-processing for multiple str's, the gap becomes even larger despite using vmap in JAX. (around 50x)

Could the people here provide any insight into how to improve the performance of the following code? (The following code is somewhat simplified from the original)

# dynamic_updates.py

@jit
def loop(
    str_arr: Array, # Array with uint8
    prefix: Array,  # initalized Array of float32
    reference: Array, #Array of float32
) -> Array | Any:
    (n,) = str_arr.shape

    def update_prefix(carry: Array, offset: int | Array) -> tuple[Array, None]:
        pref = carry
        i = jnp.arange(n, dtype=jnp.int32)
        j = i + offset

        adjacent = pref[i + 1, j - 1]
        score = reference[str_arr[i], str_arr[j], str_arr[j - 1], str_arr[i + 1]]
        adjacent_score = score + adjacent

        end_at_j = jax.scipy.special.logsumexp(
            pref[i, :-1] + pref[1:, j].T,
            axis=1,
        )
        # some more update logic here, using `log1p` and more `jax.scipy.special.logsumexp`'s 

        update_at_ij = jnp.logaddexp(adjacent_score, end_at_j)
        pref = pref.at[i, j].set(update_at_ij, mode="drop")

        return pref, None

    diag_offsets = jnp.arange(n, dtype=jnp.int32)
    new_prefix, _ = lax.scan(
        update_prefix,
        prefix,
        diag_offsets,
    )

    return new_prefix  # type: ignore

Here were some of my initial thoughts.

1. Is jsp.logsumexp very slow?
2. Are large in-place updates slowing down the code?
3. Is large dynamic slicing slowing down the code?

However, I am not exactly sure how to profile and check these thoughts either, as I would need to have a control which does not implement the above ideas to test this, and making such a control is non-trivial. Any help would be appreciated!

[bmaxdk]

Can you try with lax.fori_loop instead of lax.scan? As long as you don't need the scan carry to store the entire prefix array for each iter, then lax.fori_loop may reduce overhead:

Check

new_prefix, _ = lax.scan(
    update_prefix,
    prefix,
    jnp.arange(n),
)

to

new_prefix = lax.fori_loop(
    lower=0,
    upper=n,
    body_fun=lambda offset, pref: update_prefix(pref, offset),
    init_val=prefix,
)

This may help some improvement.

You also mentioned using vmap over multiple str_arr inputs. That will help but only each DP is large enough to get good GPU utilization and the dimension along which you vmp is large enough.