AssertionError when computing gradient of a function with reduce_window

[LeonardoTredese]

When running the following example

import jax

x = jax.numpy.ones((1, 28, 28))

def f(x):
    return jax.lax.reduce_window(x, 0., jax.lax.max, (1, 2, 2), (1, 2, 2), padding="SAME").sum()

jax.grad(f)(x)

I get the following assertion error:

traceback (most recent call last):
  File "/.... /error.py", line 8, in <module>
    jax.grad(f)(x)
  File "/..../jax/_src/traceback_util.py", line 180, in reraise_with_filtered_traceback
    return fun(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^
  File "/..../jax/_src/api.py", line 392, in grad_f
    _, g = value_and_grad_f(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/..../jax/_src/traceback_util.py", line 180, in reraise_with_filtered_traceback
    return fun(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^
  File "/..../jax/_src/api.py", line 462, in value_and_grad_f
    ans, vjp_py = _vjp(f_partial, *dyn_args)
                  ^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/..../jax/_src/api.py", line 1962, in _vjp
    out_primals, vjp = ad.vjp(flat_fun, primals_flat)
                       ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/..../jax/_src/interpreters/ad.py", line 142, in vjp
    out_primals, pvals, jaxpr, consts = linearize(traceable, *primals)
                                        ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/..../jax/_src/interpreters/ad.py", line 133, in linearize
    assert all(out_primal_pval.is_known() for out_primal_pval in out_primals_pvals)
AssertionError

I am using Python 3.11.6 and ubuntu 23.10.

jax==0.4.35
jax-cuda12-pjrt==0.4.35
jax-cuda12-plugin==0.4.35
jaxlib==0.4.35

I believe this is not an expected behaviour. Am I wrong? How can I get the gradients of this function?

[bmaxdk]

Based on your error, you may noticed jax.lax.reduce_window with jax.lax.max as the reduction function does not support reverse-mode automatic differentiation. I agree. it's odd JAX does not support gradients of lax.reduce_window with lax.max
I prefer you to use jax.nn.max_pool

import jax

x = jax.numpy.ones((1, 28, 28))

def f(x):
    return jax.nn.max_pool(x, (2, 2), (2, 2), padding="SAME").sum()

gradient = jax.grad(f)(x)
print(gradient)

[LeonardoTredese]

In the version I am using max_pool is not part of the jax.nn module.
I implemented a work around for 2d pooling on a square window, which was my usecase.

def pool2d(x, op, pool, stride):
    out_idx = jnp.arange((x.shape[-1] - pool) // stride + 1)
    def window_f(x, i, j):
        return op(jax.lax.dynamic_slice(x, (i*pool, j*pool), (pool, pool)))
    def pool_f(x):
        return jax.vmap(lambda i: (jax.vmap(lambda j: window_f(x, i, j))(out_idx)))(out_idx)
    return jax.vmap(jax.vmap(pool_f))(x)