Max pooling with reduce_window causes "ValueError: Linearization failed ..." only on some cases

[yippiez]

Using reduce_window with 0.0 as it's init value causes a Linearization error when only jit compiling a grad transformed version of the function. Below is the example code with error. The error doesn't happen if -jnp.inf is used or the current version is not jit compiled.

def maxpool_step(x, kernel_size, stride):
    return jax.lax.reduce_window(
        x,
        jnp.float32(0.0),
        jax.lax.max,
        window_dimensions=(1, 1, kernel_size, kernel_size),
        window_strides=(1, 1, stride, stride),
        padding="VALID",
    )

_grad = jax.grad(lambda x: jnp.sum(maxpool_step(x, kernel_size=2, stride=2)))
_grad_jit = jit(_grad) # => ValueError: Linearization failed ...

Interestingly removing jit compilation from the above make it not throw out this error even though it's the same. Below is the working version in which both gradand jit functions work.

def maxpool_step(x, kernel_size, stride):
    return jax.lax.reduce_window(
        x,
        -jnp.inf, #  0.0, jnp.float32(0.0), jnp.array(0.0) or jnp.array(-jnp.inf) all cause errors
        jax.lax.max,
        window_dimensions=(1, 1, kernel_size, kernel_size),
        window_strides=(1, 1, stride, stride),
        padding="VALID",
    )

My questions are:

1. Why does it throw an error for just the certain values? ( 0.0 / jnp.float32(0.0) / jnp.array(0.0) / jnp.array(-jnp.inf) )
2. In the first case why does it throw an error only when the grad function is jit compiled?