Ikeda Map in JAX

import jax
import jax.numpy as jnp
from jax import jit
from jax import lax

key = jax.random.key(0)
size = 10**6
points = jax.random.uniform(key, shape=(size, 2), minval=0, maxval=0.5)

@jit
def ikeda_update(point):
    u = 0.9
    x, y = point[0], point[1]
    t = 0.4 - (6 / (1 + x ** 2 + y ** 2))
    x_new = 1 + u * (x * jnp.cos(t) - y * jnp.sin(t))
    y_new = u * (x * jnp.sin(t) + y * jnp.cos(t))
    return jnp.array((x_new, y_new))

ikeda_update = jit(jax.vmap(ikeda_update))

@jit
def iterate(points):
    def body_fn(i, val):
        return ikeda_update(val)
    return lax.fori_loop(0, 50, body_fn, points)

points = iterate(points)
points

%%timeit
ikeda_update(points)