@louisfd The bug in WGPU about the boolean casting to float/int is no longer a blocker in the updated code. I noticed that I can get rid of that by just changing the logic a bit and adding another bool_not op which converts the weird large values of 1065353216 for (1.0 f / 1 ) to exact zeros. There's no overflow now since this op happens before the final sum, so it's all ok.

TL;DR There's no blocker for this PR and it works for all backends.

Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 78.93%. Comparing base (4427768) to head (438a7c0).
Report is 3 commits behind head on main.

@@            Coverage Diff             @@
##             main    #1342       /-   ##
==========================================
  Coverage   78.77%   78.93%    0.15%     
==========================================
  Files         551      563       12     
  Lines       61836    62981     1145     
==========================================
  Hits        48712    49714     1002     
- Misses      13124    13267      143     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

@ashdtu While we're here, can you change this LSTM test line in lstm.rs (which was the original reason I needed the any op):

 // Asserts the gradients exist and are non zero
 assert!(*some_gradient.abs().sum().into_data().value.first().unwrap() > 0.);