Help with re writing a JAX function

[hanzy1110]

Hey! I have the following function which I've been banging my head over the last several days:

def aeran_model(n_i: Array, N_i: Array, sigma_i: Array) -> Array:
    Di = 0.0
    DiPrev = 0.0

    for i, (ni, Ni, s_i) in enumerate(zip(n_i, N_i, sigma_i)):
        lnNi = jnp.log(Ni)
        delta_i = 1.25 / lnNi
        mu_i = (sigma_i[i - 1] / s_i) ** 2.0
        nip1_tot = ni
        if i > 0:
            exp_ = mu_i / delta_i
            inner = jnp.power(jnp.abs(1.0 - DiPrev), exp_)
            nip1_eff = (1.0 - inner) * Ni
            nip1_tot = nip1_eff + ni

        lin_dmg = jnp.clip(nip1_tot / Ni, 0, 1)
        inner_d = 1 - jnp.power((1 - lin_dmg), delta_i)
        DiPrev = Di
        Di = inner_d
    return jnp.abs(Di)

I then intent to vmap over the last 2 arguments. This implementation works but it seems to be quite slow, my guess is that the for loop is the reason why it is slow. Is there a way to refactor it such that it eliminates the for loop? (I've tried lax.fori_loop with no success whatsoever)
Thanks!

[jakevdp]

Hi - it looks like you could replace this with a scan, but it would be helpful if you could provide a complete example of expected inputs. Are these one-dimensional arrays? I tried the following:

n_i = np.random.randn(10)
N_i = np.random.randn(10)
sigma_i = np.random.randn(10)

And the output is NaN. Is that expected?