import jax.ops

index = jax.ops.index
index_update = jax.ops.index_update
index_add = jax.ops.index_add
index_mul = jax.ops.index_mul

import jax.ops

index = jax.ops.index
index_update = jax.ops.index_update
index_add = jax.ops.index_add
index_mul = jax.ops.index_mul
