Example

softmax_axis1.kina

Source

tensor<F32, [2, 3], @host>.softmax(axis=1)