Implementing L2L in JAX

[davisyoshida]

I'm interested in implementing a decorator to swap parameters on and off of the GPU as described here: https://arxiv.org/pdf/2002.05645.pdf

To get started, I made the following decorator:

  _cpu = jax.devices('cpu')[0]
  
  def swap_params(func):
      @jax.custom_vjp
      def inner(params, *args):
          gpu_params = jax.tree_map(lambda v: jax.device_put(v, _gpu), params)
          return func(gpu_params, *args)
  
      def fwd(params, *args):
          result = inner(params, *args)
          return result, (params, args)
  
      def bwd(res, tangents):
          params, args = res 
          gpu_params = jax.tree_map(lambda v: jax.device_put(v, _gpu), params)
          param_grad, *arg_grad =  jax.vjp(func, gpu_params, *args)[1](tangents)
          param_grad = jax.device_put(param_grad, _cpu)
          return param_grad, *arg_grad
  
      inner.defvjp(fwd, bwd)
      return inner

Unfortunately, functions decorated by this cannot be JIT-ed, since they take arguments living on both the CPU and GPU. Is there a way I can do something like this (or some other method entirely) and still make use of jax.jit?