pub fn softmax<T: Float>(x: &Tensor<T>, _dim: i64) -> Result<Tensor<T>>
Applies softmax along the specified dimension.