Function checknonfinite_lower

imperative/python/megengine/xla/rules/math.py:560–567

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

@register_lower_rule(mops.CheckNonFinite)
def checknonfinite_lower(ctx, *args: Union[HLOTensor, Sequence[HLOTensor]]):
    check_inf = [logical_or(isnan(arg), isinf(arg)) for arg in args]
    check_inf_reduce = [expand_dims(sum(inp), 0) for inp in check_inf]
    concat_res = concat(check_inf_reduce, 0)
    if_nonfinite = expand_dims(sum(concat_res), 0).astype(np.int32)
    res = [arg * ctx.op.scale for arg in args]
    res.append(if_nonfinite)
    return res