pub trait SoftmaxAxis: Softmax { // Required method fn softmax_axis(self, axis: usize) -> Self::Output; }