Function softmax

python/tvm/contrib/cudnn.py:877–900

Compute softmax using CuDNN Parameters ---------- x : tvm.te.Tensor The input tensor axis : int The axis to compute the softmax Returns ------- ret : tvm.te.Tensor The result tensor

(x, axis=-1)

def softmax(x, axis=-1):
    """Compute softmax using CuDNN

    Parameters
    ----------
    x : tvm.te.Tensor
        The input tensor

    axis : int
        The axis to compute the softmax

    Returns
    -------
    ret : tvm.te.Tensor
        The result tensor
    """
    return te.extern(
        x.shape,
        [x],
        lambda ins, outs: tvm.tirx.call_packed(
            "tvm.contrib.cudnn.softmax.forward", ins[0], outs[0], axis
        ),
        name="y",
    )


def log_softmax(x, axis=-1):