jax.numpy.nanstd bad behaviour with jax.grad?

[eserie]

Hi,

I'm trying to compute gradient of a function that run on data which contains nan values and use jnp.nanstd.

I run the following experiment:

rng = jax.random.PRNGKey(42)
x = jax.random.normal(rng, (10,))
_, rng = jax.random.split(rng)
params = {"w": jax.random.normal(rng, (10,))}

# put some nan values
x = jax.ops.index_update(x, 0, jnp.nan)

Then run:

def test(params, x):
    y = params["w"] * x
    score = jnp.nanstd(y)
    return score

jax.value_and_grad(test)(params, x)

This gives:

(DeviceArray(0.3482654, dtype=float32),
 {'w': DeviceArray([nan, nan, nan, nan, nan, nan, nan, nan, nan, nan], dtype=float32)})

However:

def test(params, x):
    y = params["w"] * x
    
    mean = jnp.nanmean(y)
    mean2 = jnp.nanmean(y ** 2)
    score = jnp.sqrt(mean2 - mean ** 2)
    return score

jax.value_and_grad(test)(params, x)

gives:

(DeviceArray(0.3482654, dtype=float32),
 {'w': DeviceArray([        nan,  0.01182435, -0.06609891,  0.21879005,
                0.09108952,  0.06071524, -0.00908424,  0.0750377 ,
                0.01525824, -0.0830379 ], dtype=float32)})

Finally, the third way to do:

def test(params, x):
    y = params["w"] * x

    mean = jnp.nanmean(y)
    var = jnp.nanmean((y - mean) ** 2)
    score = jnp.sqrt(var)
    return score


jax.value_and_grad(test)(params, x)

gives:

(DeviceArray(0.3482654, dtype=float32),
 {'w': DeviceArray([nan, nan, nan, nan, nan, nan, nan, nan, nan, nan], dtype=float32)})

Is this a bug in jnp.nanstd ? Do you have some recommendations for such use cases?

[jakevdp]

Thanks for the report – I think this is a bug, so I'm going to convert this discussion to an issue: see #8128.

[eserie]

Thanks for taking a look !