donate_argnums and CPU

[davisyoshida]

I wasn't sure if I should open an issue about this or if this is more of an XLA problem but it seems that passing donate_argnums to jit when using CPU doesn't work:

import jax 
import jax.numpy as jnp 

for dev_name in ('cpu', 'gpu'):
    print(f'{dev_name}:')
    device = jax.devices(dev_name)[0]
    func = jax.jit(jnp.square, donate_argnums=(0,), device=device)
    func(jax.device_put(jnp.ones((5, 5)), device)).block_until_ready()

# Output:
# cpu:
# .../xla.py:708: UserWarning: Some donated buffers were not usable: f32[2,2]{1,0}
# [[1. 1.]
#  [1. 1.]]
# gpu:
# [[1. 1.]
#  [1. 1.]]

I was looking into doing some parameter updates on the CPU and was surprised to find that JITing jnp.add didn't allow me to donate the parameter buffer. Is there any way I can force the JIT to write the result to the input buffer to avoid a copy?