# Softmax in Rust and Nalgebra Style

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 a python example as reference:

``````def softmax(x):
return np.exp(x) / np.sum(np.exp(x), axis=0)
``````

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);
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)` 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])
}
``````
1 Like

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.

1 Like