TensorF32.softmax: def of stelliferous@0.0.2.0
Import: import stelliferous@0.0.2.0/tensor_f32.bend as TensorF32
Lines 383-390 of stelliferous@0.0.2.0/tensor_f32.bend: https://jend.leeeet.dev/read.txt?path=stelliferous@0.0.2.0/tensor_f32.bend&start_line=383&end_line=390

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

# Keep a singleton reduction axis for broadcasting. Both reductions use the
# existing ordered fold; reduction gathers row bands to a dense owner.
def softmax(tensor: Tensor(),+axis: U32) -> Tensor():
  then_shape(shape(tensor),tensor => extents =>
    +unit = Layout.with_extent(extents,U32.to_nat(axis),1)
    then(copy(tensor),values => maximum_input =>
      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)))),
        values => sum_input => div(values,reshape(Layout.reduce_axis(~F32,~F32.add,"softmax",0.0,sum_input,axis),unit)))))

Uses (18):
  https://jend.leeeet.dev/definition.txt?path=Base/base.bend&name=F32
  https://jend.leeeet.dev/definition.txt?path=Base/base.bend&name=F32.add
  https://jend.leeeet.dev/definition.txt?path=Base/base.bend&name=F32.div
  https://jend.leeeet.dev/definition.txt?path=Base/base.bend&name=F32.max
  https://jend.leeeet.dev/definition.txt?path=Base/base.bend&name=F32.neg
  https://jend.leeeet.dev/definition.txt?path=Base/base.bend&name=U32
  https://jend.leeeet.dev/definition.txt?path=Base/base.bend&name=U32.to_nat
  https://jend.leeeet.dev/definition.txt?path=stelliferous@0.0.2.0/tensor_f32.bend&name=Tensor
  https://jend.leeeet.dev/definition.txt?path=stelliferous@0.0.2.0/tensor_f32.bend&name=copy
  https://jend.leeeet.dev/definition.txt?path=stelliferous@0.0.2.0/tensor_f32.bend&name=div
  https://jend.leeeet.dev/definition.txt?path=stelliferous@0.0.2.0/tensor_f32.bend&name=exp
  https://jend.leeeet.dev/definition.txt?path=stelliferous@0.0.2.0/tensor_f32.bend&name=reshape
  https://jend.leeeet.dev/definition.txt?path=stelliferous@0.0.2.0/tensor_f32.bend&name=shape
  https://jend.leeeet.dev/definition.txt?path=stelliferous@0.0.2.0/tensor_f32.bend&name=sub
  https://jend.leeeet.dev/definition.txt?path=stelliferous@0.0.2.0/tensor_f32.bend&name=then
  https://jend.leeeet.dev/definition.txt?path=stelliferous@0.0.2.0/tensor_f32.bend&name=then_shape
  https://jend.leeeet.dev/definition.txt?path=stelliferous@0.0.2.0/tensor_layout.bend&name=reduce_axis
  https://jend.leeeet.dev/definition.txt?path=stelliferous@0.0.2.0/tensor_layout.bend&name=with_extent