dynamic_slice_update example

[AdrienCorenflos]

Hi,
I'm trying to fill forward an array in the following fashion

while_loop...
arr = index_update(arr, index[i + 1: k], arr[i])

JAX is complaining about it due to the need for the program to be static, and is telling me to use dynamic_update_slice instead. Is this something I can do or am I condemned to update element by element? Note that i and k are dynamic.

[jakevdp]

dynamic_slice_update can only be used for operations where the slice size is static (but the start location is dynamic); for example:

from jax import lax
import jax.numpy as jnp
x = jnp.zeros(10)
y = jnp.ones(3)
lax.dynamic_update_slice(x, y, (3,))
# DeviceArray([0., 0., 0., 1., 1., 1., 0., 0., 0., 0.], dtype=float32)

One way to achieve the operation you have in mind is using jnp.where:

ind = jnp.arange(where.shape[0])
arr = jnp.where((ind > i) & (ind < k), arr[i], arr)

[AdrienCorenflos]

Cool thanks!

[jakevdp]

The docstring lacked any examples, so I've added a few in #6550

[AdrienCorenflos]

👍 Le ven. 23 avr. 2021 à 19:12, Jake Vanderplas ***@***.***> a écrit :

…

The docstring lacked any examples, so I've added a few in #6550 <#6550> — You are receiving this because you authored the thread. Reply to this email directly, view it on GitHub <#6547 (reply in thread)>, or unsubscribe <https://github.com/notifications/unsubscribe-auth/AEYGFZ47GXF27T5W6A6HENTTKGL7FANCNFSM43OG2IKQ> .