Kina.ai

Example

softmax_axis1.kina

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