Explicit managing the output buffer of `jax.jit` function about jax HOT 1 OPEN

from functools import partial
import jax
import jax.numpy as jnp

state = (jnp.zeros(4), jnp.arange(4))

@partial(jax.jit, donate_argnums=0)
def f(state):
  x, y = state
  x += 1
  y *= 2
  return (x, y)

pointer0 = state[0].unsafe_buffer_pointer()
pointer1 = state[1].unsafe_buffer_pointer()

state = f(state)

assert state[0].unsafe_buffer_pointer() == pointer0  # same memory
assert state[1].unsafe_buffer_pointer() == pointer1  # same memory