Overide donate_argnums

[pharringtonp19]

Is it possible to tell jax at the top of a script to ignore all donate_argnum statements?

Working with an optimizer provided by KFAC-JAX, I'm getting the following error.

RuntimeError: INVALID_ARGUMENT: Invalid buffer passed to Execute() as argument 0 to replica 0: INVALID_ARGUMENT: Buffer has been deleted or donated.

[YouJiacheng]

For a temporary fix, you can modify the source code of JAX.

[jakevdp]

It looks like kfac_jax consumes its input buffers. You could work around it like this:

args = (xs, ys)
args2 = (xs.copy(), ys.copy())  # necessary because kfac_jax donates the input buffers
print(jax.value_and_grad(loss_fn)(params, args))
print(jax.value_and_grad(loss_fn_inner_step)(params, jax.random.PRNGKey(0), args2))

This might be unintended on the part of kfac_jax; I'd probably raise an issue on the kfac_jax repo

[pharringtonp19]

@YouJiacheng Thanks for the response.

I think the only thing we cannot do in my particular case is apply jit to the outer function f (in my example: loss_fn_inner_step). When I do, I get a ConcretizationTypeError related to the verify_args_and_get_step_counter method in kfac-jax.

import jax 
import jax.numpy as jnp
import kfac_jax

def model(params, x):
  return params['w']*x + params['b']

def sq_error(predictions: jnp.ndarray, targets: jnp.ndarray):
  kfac_jax.register_squared_error_loss(predictions, targets)
  return jnp.mean((predictions - targets)**2)

def loss_fn(params, batch):
  x, y = batch
  yhat = jax.vmap(model, in_axes=(None,0))(params, x)
  return sq_error(yhat, y)

def inner_step(params, key, batch):
  subkey1, subkey2= jax.random.split(key)
  opt_state = optimizer.init(params, subkey1, batch)
  params, _, _ = optimizer.step(params, opt_state, subkey2, batch=batch)
  return params

@jax.jit
def loss_fn_inner_step(params, rng, batch):
  with jax.disable_jit():
    params = inner_step(params, rng, batch)
  return loss_fn(params, batch)

optimizer = kfac_jax.Optimizer(
    value_and_grad_func=jax.value_and_grad(loss_fn),
    l2_reg=0,
    value_func_has_aux=False, 
    value_func_has_state=False,
    value_func_has_rng=False, 
    use_adaptive_learning_rate=True,
    use_adaptive_momentum=True,
    use_adaptive_damping=True,
    initial_damping=1.0,
    multi_device=False,
)

params = {"w": jax.random.normal(jax.random.PRNGKey(0), shape=((1,))), "b": jax.random.normal(jax.random.PRNGKey(1), shape=((1,)))}
xs = jnp.linspace(-1,1,100).reshape(-1,1)
ys = jnp.sin(xs)

print(jax.value_and_grad(loss_fn)(params, (xs, ys)))
print(jax.value_and_grad(loss_fn_inner_step)(params, jax.random.PRNGKey(0), (xs, ys)))

ConcretizationTypeError                   Traceback (most recent call last)
[/usr/local/lib/python3.7/dist-packages/kfac_jax/_src/optimizer.py](https://localhost:8080/#) in verify_args_and_get_step_counter(self, step_counter, learning_rate, momentum, damping, global_step_int)

[jakevdp]

(wanted to make sure you also saw my response above, which is I think a cleaner fix)

[pharringtonp19]

@jakevdp Thanks for the help! Based on my current understanding of python/jax I am surprised that copy() made a difference. Could you point me towards a resource that explains how this works? --thanks

[YouJiacheng]

@pharringtonp19 donate_argnums donate the buffer, while copy copy the data and the buffer. i.e. for x_copy = x.copy() x's buffer can be donated without affecting x_copy.