fix(transforms): reset notfinite_count after apply_if_finite gives up and accepts non-finite update

[wahid18-maqs]

Summary

This PR fixes the apply_if_finite state tracking by resetting notfinite_count after the optimizer gives up and accepts a non-finite update.

Changes

• Reset notfinite_count to 0 on the step where apply_if_finite exceeds max_consecutive_errors and accepts the update.
• Add a regression test verifying the counter is reset after the give-up step.
• Extend the regression to cover the pmap path as well.

Why

Previously, notfinite_count continued to increase after apply_if_finite exceeded max_consecutive_errors and accepted a non-finite update. This caused the counter to represent the total accumulated non-finite updates rather than the current streak of consecutive failures.

Resetting the counter after the give-up step makes its behavior consistent with its intended semantics and provides more meaningful diagnostic information during training.