I’m new to rust and just started diving in to nalgebra this weekend. I wrote a naive version of softmax to get more familiar with nalgebra. Curious if anyone has advice in whether my implementation could be better?

Here is my nalgebra example using just a 4x4 matrix:

fn softmax(m: &mut Matrix4<f64>) {
m.apply(|v| v.exp());
let mut v4 = Vector4::zeros();
let (r, c) = m.shape();
for idx in 0..c {
let cm = m.column_mut(idx);
let sum = cm.iter().fold(0.0, |sum, val| sum + val);
v4[idx] = sum;
}
for rw in 0..r {
let mut rv = RowVector4::zeros();
for (col, val) in v4.iter().enumerate() {
let ij = m.row(rw).column(col)[0];
rv[col] = ij / val;
}
m.set_row(rw, &rv);
}
}

Here are a few alternatives. First, your code can be simplified at some places:

You can replace the fold by a sum.

You can access a specific component of a matrix using indexing. Thus you can replace the m.row(rw).column(col)[0] by m[(rw, col)].

In your second loop, you may just modify m in-place with m[(rw, col)] /= val.

m.apply(|v| v.exp());
let mut v4 = Vector4::zeros();
let (r, c) = m.shape();
for idx in 0..c {
let cm = m.column(idx);
let sum: f64 = cm.iter().sum();
v4[idx] = sum;
}
for rw in 0..r {
for (col, val) in v4.iter().enumerate() {
m[(rw, col)] /= val
}
}

Then, the first for loop can be compacted in a single line, using the from_fn constructor, and the inner loop on the second for can be replaced by a component-wise division. Note that because dimensions must match, the v4 is now a RowVector and we use the column index (the second one) passed to the callback given to the ::from_fn constructor:

m.apply(|v| v.exp());
let v4 = RowVector4::<f64>::from_fn(|_, j| m.column(j).iter().sum());
for rw in 0..m.nrows() {
m.row_mut(rw).component_div_assign(&v4);
}

Finally, iterating on rows can be much slower than iterating on columns if you intend to implement this for larger matrices. Here is a version that iterate on each column:

m.apply(|v| v.exp());
let v4 = Vector4::<f64>::from_fn(|i, _| m.column(i).iter().sum());
for j in 0..m.ncols() {
let mut c = m.column_mut(j);
c /= v4[j];
}

Alternatively, if you use std::ops::DivAssign this can be written in a slightly more compact way:

m.apply(|v| v.exp());
let v4 = Vector4::<f64>::from_fn(|i, _| m.column(i).iter().sum());
for j in 0..m.ncols() {
m.column_mut(j).div_assign(v4[j])
}

wow, nice. thanks for responding. i knew there had to be a better way. I’m looking to write up a blog post on rust, nalgebra, and softmax… I’ll link back here when it’s done!

BTW, I’m really digging the nalgebra project. Does there exist (or any plans) for a book on rust and nalgebra? Or do any rust books cover nalgebra in any detail?

Great! Don’t hesitate to ping me if you need a review of your post.

Right now the main source of documentation for nalgebra is the user guide on https://nalgebra.org.
No existing book I am aware of covers nalgebra in details. As far as I’m concerned I don’t plan on writing a book in the near future by lack of time unfortunately. I will consider writing one if I manage to work full-time on my rust projects at some point in the future.