I pushed changes for the backward and gradient pass. While debugging those, I realized that the forward pass on the softmax didn't work correctly when the masks used the intended values of -Inf. So, I implemented a custom forward masked softmax pass manually. It's not particularly pretty but it works correctly.

I believe there are a ton of optimizations that can be made to this implementation (mostly in memory allocations), but we can return to that later---correctness is the important part here.

Using std::numeric_limits::lowest() is actually risky, I just realized. When I was testing the forward pass locally, I set the attention mask to all zeros, so I didn’t notice this problem at the time.
Do you think there’s any way we can still use softmax.forward() later when we optimize the implementation?

Yeah, the idea was to actually use -Inf as the masking value, not -DBL_MAX.

I'm not sure if using softmax.Forward() will be possible---I think the implementation here will change a lot when it gets optimized. We'll see?