Softmax in Rust and Nalgebra Style

1

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)[0];
            rv[col] = ij / val;
        }
        m.set_row(rw, &rv);
    }
}
2

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

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?

4

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