def METHOD_NAME(softmax_op, s, outs):
    op_tag = softmax_op.tag
    if op_tag == "softmax_output":
        exp = softmax_op.input_tensors[0]
        expsum = softmax_op.input_tensors[1]
        max_elem = s[exp].op.input_tensors[1]
        delta = None
        axis = int(softmax_op.attrs["axis"])
    elif op_tag == "fast_softmax_output":
        exp = softmax_op.input_tensors[0]
        expsum = softmax_op.input_tensors[1]
        delta = s[exp].op.input_tensors[0]
        max_elem = s[delta].op.input_tensors[1]
        axis = int(softmax_op.attrs["axis"])
    elif op_tag == "log_softmax_output":
        exp = None
        delta = None
        max_elem = softmax_op.input_tensors[1]
        expsum = softmax_op.input_tensors[2]
        axis = int(softmax_op.attrs["axis"])
    else:
        raise ValueError(
            f"Tag is expected to be softmax_output or log_softmax_output. Got {op_tag}"
        )
    output = outs[0]
    def _schedule(output_op, softmax_op):
        # only parallelize outer dimensions up to axis
        outer_axes = [output_op.axis[i] for i in range(0, axis)]
        fused_outer_axes = s[output_op].fuse(*outer_axes)
        s[output_op].parallel(fused_outer_axes)
        if softmax_op != output_op:
            # fuse softmax output with following elemwise ops.
            s[softmax_op].compute_at(s[output_op], fused_outer_axes)
        # move computations with the same outer dimensions under the same root
        s[max_elem].compute_at(s[output_op], fused_outer_axes)
        s[expsum].compute_at(s[output_op], fused_outer_axes)
        if delta is not None:
            s[exp].compute_inline()
            s[delta].compute_inline()
        if exp is not None:
            s[exp].compute_at(s[output_op], fused_outer_axes)
    if list(output.shape) == list(softmax_op.output(0).shape):
        _schedule(output.op, softmax_op)
    else:
        # This case can happen, for example, if the 4D input to softmax
        # is in the NCHW layout while the fused elemwise op takes the NCHWc layout.
        # Since we parallelize over outer axes up to the "axis" parameter of softmax,
        # softmax and the fused op need to be in the same layout if we want to
        # fuse them under the same parallel loop.
        # This case can be removed if softmax supported AlterLayout.
        schedule_injective_from_existing(s, output)
        _schedule(softmax_op, softmax_op)