def METHOD_NAME():
    if tgt.kind.name in ["nvptx", "rocm"]:
        dtype = softmax_op.output(0).dtype
        return dtype in ["float32", "int32"]
    if tgt.kind.name != "cuda":
        # this is used as the gpu schedule for other arches which
        # may not have warp reductions
        return False
    return True