Boolean indexing alternative

[Kheims]

Hi, I have seen that boolean indexing isn't yet implemented.

Particularly, I have seen @awni's workaround recommandation ( #246) on using mx.where, but apparently it has issues handling inf values for x (it's a TODO in the cpp operation), which is the case for us.

Here is our usecase :

        if normalize:
            deg = self._degree(col, x.shape[0])
            deg_inv_sqrt = deg ** (-0.5)
            # NOTE : need boolean indexing in order to zero out inf values 
            deg_inv_sqrt[deg_inv_sqrt == float('inf')] = 0 # doesn't work yet
            deg_inv_sqrt = mx.where(deg_inv_sqrt == float('inf'), mx.zeros_like(deg_inv_sqrt), deg_inv_sqrt) # Nan in place of zeros for inf values
            norm = deg_inv_sqrt[row] * deg_inv_sqrt[col]

Is there an alternative to using numpy or worse a listcomp ?

[awni]

That's a tricky one 🤔 .

Of course we need to fix the where to handle the inf case. That is not conceptually difficult, but might require a primitive so it's a bit of work.

Just curious why does deg have inf in it? Is it a sentinal for something? If so could you use a different value? Like maybe just something very large?

[Kheims]

Both solutions didn't solve the issue unfortunately :(.

a = mx.array([1,2,3, float('inf')])
mx.where(a < float('inf'), a, 0) # output : [1,2,3, nan]

[awni]

You're right disregard those, sorry they are not correct, the NaNs just get moved..

[Kheims]

I guess the viable solution for now is to create the indices with numpy ? But we for sure will need a consistent solution as this will be an important step in GCNConv's forward.
If you can point me to the right direction i'd gladly do my best to contribute to the where operation.

[awni]

@Kheims It shouldn't be too bad to precompute the non-zero degree indices from a perf standpoint. I know it is not that elegant.

A slightly hacky solution could be to do something like:

>>> import mlx.core as mx
>>> deg = mx.array([1,2,3,0])
>>> deg += 1e-10
>>> deg_inv_sq = deg**(-0.5)
>>> deg_inv_sq = mx.where(deg_inv_sq <= 1, deg_inv_sq, 0)
>>> deg_inv_sq
array([1, 0.707107, 0.57735, 0], dtype=float32)

[Kheims]

Thank you for the workaround :)