def softmax

stelliferous@0.0.2.0 · tensor_f32.bend:383

import Base
import stelliferous@0.0.2.0/tensor_layout.bend as Layout

383# Keep a singleton reduction axis for broadcasting. Both reductions use the
384# existing ordered fold; reduction gathers row bands to a dense owner.
385def softmax(tensor: Tensor(),+axis: U32) -> Tensor():
386  then_shape(shape(tensor),tensor => extents =>
387    +unit = Layout.with_extent(extents,U32.to_nat(axis),1)
388    then(copy(tensor),values => maximum_input =>
389      then(copy(exp(sub(values,reshape(Layout.reduce_axis(~F32,~F32.max,"softmax",F32.div(F32.neg(1.0),0.0),maximum_input,axis),unit)))),
390        values => sum_input => div(values,reshape(Layout.reduce_axis(~F32,~F32.add,"softmax",0.0,sum_input,axis),unit)))))