jax.remat with static_argnums

[clemisch]

Dear jax-Team,

I want to use jax.remat on a function which has static_argnums. The static argument is another function. When I compute the gradient, JAX complains about the function being not a valid JAX type. Is there a way to mark an argument as static in jax.remat, similar to jax.jit?

---

import jax
import jax.numpy as jnp

@jax.remat
def fun(x, f):
    y = f(x)
    return jnp.sum(jnp.square(y))

fun_grad = jax.jit(jax.grad(fun), static_argnums=1)     # same without JIT

x = arange(10.)
fun_grad(x, jnp.sin)

gives

TypeError                                 Traceback (most recent call last)
<ipython-input-16-4b715c26497f> in <module>
     10 fun_grad = jax.grad(fun)
     13 x = arange(10.)
---> 14 fun_grad(x, jnp.sin)

    [... skipping hidden 18 frame]

~/.local/lib/python3.8/site-packages/jax/core.py in concrete_aval(x)
    979 if hasattr(x, '__jax_array__'):
    980   return concrete_aval(x.__jax_array__())
--> 981 raise TypeError(f"Value {repr(x)} with type {type(x)} is not a valid JAX "
    982                  "type")

TypeError: Value <CompiledFunction of <function _one_to_one_unop.<locals>.<lambda> at 0x7ff252ffd820>> with type <class 'CompiledFunction'> is not a valid JAX type

[niklasschmitz]

A late reply, but just in case here's the fix:
jax.remat got subsumed by jax.checkpoint which also supports static_argnums:

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

@partial(jax.checkpoint, static_argnums=1) # new jax.checkpoint API supports static_argnums
def fun(x, f):
    y = f(x)
    return jnp.sum(jnp.square(y))

fun_grad = jax.jit(jax.grad(fun), static_argnums=1)  

x = jnp.arange(10.)
fun_grad(x, jnp.sin)

DeviceArray([ 0.        ,  0.90929735, -0.7568025 , -0.2794155 ,
              0.98935825, -0.5440211 , -0.5365729 ,  0.9906074 ,
             -0.2879033 , -0.75098723], dtype=float32)