Is there a JAX function like torch.Tensor.masked_fill_?

[ayaka14732]

torch.Tensor.masked_fill_

>>> import numpy as np
>>> import torch
>>> a = torch.rand((3, 2, 5))
>>> a
tensor([[[0.6125, 0.2513, 0.4099, 0.8357, 0.6188],
         [0.0068, 0.5963, 0.3633, 0.4969, 0.2056]],

        [[0.7309, 0.3909, 0.8121, 0.9801, 0.0426],
         [0.4047, 0.1876, 0.5262, 0.4997, 0.4828]],

        [[0.5044, 0.8133, 0.0537, 0.4708, 0.0057],
         [0.6625, 0.6482, 0.2043, 0.5460, 0.2341]]])
>>> mask = torch.tensor([[True, True, True, False, False],
...                      [True, True, False, False, False]])
>>> a.masked_fill_(mask, np.NINF)
tensor([[[  -inf,   -inf,   -inf, 0.8357, 0.6188],
         [  -inf,   -inf, 0.3633, 0.4969, 0.2056]],

        [[  -inf,   -inf,   -inf, 0.9801, 0.0426],
         [  -inf,   -inf, 0.5262, 0.4997, 0.4828]],

        [[  -inf,   -inf,   -inf, 0.4708, 0.0057],
         [  -inf,   -inf, 0.2043, 0.5460, 0.2341]]])

[tomhennigan]

I'm not aware of a built in, but you can do this with broadcast + select:

>>> def masked_fill(mask, a, fill):
...   return jax.lax.select(mask, a, jax.lax.broadcast(fill, a.shape))

>>> key = jax.random.PRNGKey(42)
>>> key1, key2 = jax.random.split(key)
>>> a = jax.random.normal(key1, (4, 4))
>>> mask = jax.random.bernoulli(key2, shape=a.shape)
>>> masked_fill(mask, a, jnp.inf)
DeviceArray([[        inf,         inf, -1.9382184 , -0.9676806 ],
             [-0.3920406 ,  0.6062071 ,  0.37990323,         inf],
             [ 1.3282976 ,  1.1882836 ,         inf,  1.2611461 ],
             [        inf,         inf,  0.09786294,         inf]],            dtype=float32)

[ayaka14732]

Thanks! I am still expecting a built-in way to do this.

[hawkinsp]

I think jnp.where does pretty much what you want, because it has broadcasting semantics:

In [1]: import numpy as np, jax, jax.numpy as jnp

In [2]: x = jax.random.normal(jax.random.PRNGKey(42), (3, 2, 5))
WARNING:absl:No GPU/TPU found, falling back to CPU. (Set TF_CPP_MIN_LOG_LEVEL=0 and rerun for more info.)

In [3]: x
Out[3]:
DeviceArray([[[-0.735987  ,  0.11322729, -1.3516467 , -0.47371545,
                0.04204526],
              [ 0.21723728,  0.95681196,  0.14842199, -1.1293895 ,
               -1.1806034 ]],

             [[ 1.3572347 ,  2.306847  , -0.5553498 , -0.04713134,
                1.2040526 ],
              [ 1.3805705 , -2.0624716 , -1.5747548 , -0.37440643,
                1.188334  ]],

             [[-0.3921841 ,  0.38108912,  1.0760009 , -0.33681598,
               -0.96420234],
              [ 0.37908515,  0.96591413, -0.70901006, -1.085212  ,
                0.92485505]]], dtype=float32)

In [4]: mask = jnp.array([[True, True, True, False, False], [True, True, False, False, False]])

In [5]: jnp.where(mask, -jnp.inf, x)
Out[5]:
DeviceArray([[[       -inf,        -inf,        -inf, -0.47371545,
                0.04204526],
              [       -inf,        -inf,  0.14842199, -1.1293895 ,
               -1.1806034 ]],

             [[       -inf,        -inf,        -inf, -0.04713134,
                1.2040526 ],
              [       -inf,        -inf, -1.5747548 , -0.37440643,
                1.188334  ]],

             [[       -inf,        -inf,        -inf, -0.33681598,
               -0.96420234],
              [       -inf,        -inf, -0.70901006, -1.085212  ,
                0.92485505]]], dtype=float32)

[ayaka14732]

Thank you all!