Function where

imperative/python/megengine/xla/rules/tensor.py:219–224

(mask, x, y)

def where(mask, x, y):
    assert isinstance(mask, HLOTensor), f"mask must be HLOTensor, get {type(mask)}"
    x = x if isinstance(x, HLOTensor) else HLOTensor(x)
    y = y if isinstance(y, HLOTensor) else HLOTensor(y)

    return mask * x + (np.array(1.0).astype(x.dtype) - mask) * y


def where_grad(dout, mask):