x.at[].set() operation causes memcpyD2H，which is time-consuming

[Sun-Xiaohui]

Hi, JAX Team. Due to jax.array's immutability, append or insert operation in Numpy should be converted to x.at[].set() in JAX. However, it cases Memcpy, which is time-consuming. Recently, I implemented the CRF's viterbi_decode function based on JAX, but the speed is slower than Pytorch even if applying @jax.jit . Below is the main part of viterbi_decode:

def viterbi_compute_best_path(carry:tuple, batch_idx: int):
    seq_lens, score, path, end_trans = carry
    seq_end_idx = seq_lens[batch_idx] - 1
    best_last_label = (score[seq_end_idx][batch_idx] + end_trans).argmax(0).astype(int)
    best_labels = jnp.zeros(seq_len, dtype=jnp.int32)
    best_labels = best_labels.at[seq_end_idx].set(best_last_label)

    def viterbi_backward_predict_labels(i, best_labels):
      p = path[-i]
      best_last_label = p[batch_idx][best_labels[-i + 1]]
      best_labels = best_labels.at[-i].set(best_last_label)
      return best_labels

    best_labels = jax.lax.fori_loop((-seq_end_idx + 1), 1, viterbi_backward_predict_labels, best_labels)
    return carry, best_labels
carry = (seq_lens, jnp.array(score), jnp.array(path), self.end_trans) 
best_paths = jax.lax.scan(viterbi_compute_best_path, carry, jnp.arange(batch_size))[1]

I got profiling result like this by using JAX profiler:

A large number of MemcpyD2H operation are brought in. In fact, I have achieved an obvious acceleration of model training(NER and LLM) based on @jax.jit, but testing is slower than Pytorch due to x.at[].set(). How could I bypass x.at[].set() or is there any way to speed it? Thanks for your time!