x = jnp.array([1, 2, 3, 4, 5])

# Set a value
x.at[0].set(10)         # [10, 2, 3, 4, 5]

# Add to a value
x.at[0].add(10)         # [11, 2, 3, 4, 5]

# Multiply a value
x.at[0].multiply(10)    # [10, 2, 3, 4, 5]

# Works with slices too
x.at[1:3].set(99)       # [1, 99, 99, 4, 5]

        __
        __