keras/src/backend/jax/random.py/beta
def beta(shape, alpha, beta, dtype=None, seed=None):
    dtype = dtype or floatx()
    seed = jax_draw_seed(seed)
    # jax doesn't accept python lists as arguments
    alpha = jax.numpy.array(alpha)
    beta = jax.numpy.array(beta)
    sample = jax.random.beta(
        key=seed, a=alpha, b=beta, shape=shape, dtype=dtype
    )
    return sample