Function where_lower

imperative/python/megengine/xla/rules/tensor.py:786–791

(ctx, *args: Union[HLOTensor, Sequence[HLOTensor]])

@register_lower_rule(mops.Where)
def where_lower(ctx, *args: Union[HLOTensor, Sequence[HLOTensor]]):
    assert len(args) == 3 and len(ctx.vars_in) == 3 and len(ctx.vars_out) == 1
    assert _shape_equal(ctx.vars_in[0].shape, ctx.vars_in[1].shape) and _shape_equal(
        ctx.vars_in[0].shape, ctx.vars_in[2].shape
    )
    return where(args[0], args[1], args[2])


@register_lower_rule("WhereBackward", mops.WhereBackward)