Function fn

transformer_layer.py:48–49

(array: Array, idx: Array)

def slice_in_dim_1(window_length: int) -> Callable[[Array, Array], Array]:
  @jax.vmap
  def fn(array: Array, idx: Array) -> Array:
    return jax.lax.dynamic_slice_in_dim(array, idx, window_length, axis=0)

  return fn