o X۷i@s"ddlmZddlZdddZdS)) annotationsNcCs2tj||dd}t||}|tj||ddS)aSoftmax function. The softmax function transforms each element of a collection by computing the exponential of each element divided by the sum of the exponentials of all the elements. Parameters ---------- x : array-like The input array axis : int or tuple of ints, optional Axis to compute values along. Default is None Returns ------- s : cupy.ndarray Returns an array with same shape as input. The result will sum to 1 along the provided axis T)axiskeepdims)cupyamaxexpsum)xrx_max exp_x_shiftedr R/home/ubuntu/vllm_env/lib/python3.10/site-packages/cupyx/scipy/special/_softmax.pysoftmaxsr)N) __future__rrrr r r r s