def _jitted_nmapped_update(
    array: jax.Array,
    index_thunks: tuple[_StaticThunk | _DynamicThunk | _SliceThunk, ...],
    values: jax.Array,
    method: str,
    *,
    indices_are_sorted=False,
    unique_indices=False,
    mode=None,
):
  """JIT-compiled helper for in-place updates."""
  indexer = tuple(thunk.unwrap() for thunk in index_thunks)
  return getattr(array.at[indexer], method)(
      values,
      indices_are_sorted=indices_are_sorted,
      unique_indices=unique_indices,
      mode=mode,
  )
JIT-compiled helper for in-place updates.
penzai/core/named_axes.py