fix: initialize schedule_free step_count to zero

[wahid18-maqs]

Summary

This PR fixes an off-by-one issue in schedule_free by initializing step_count to 0 instead of 1.

Changes

• Initialize ScheduleFreeState.step_count to 0.
• Add a regression test verifying:
  • the initial step_count is zero, and
  • the first optimizer update evaluates the learning-rate schedule at step 0.

Why

Previously, step_count was initialized to 1 while the learning-rate schedule was queried before the counter was incremented. As a result, the first optimizer update skipped the schedule's step-0 value, leading to incorrect behavior for schedules that depend on the initial step (e.g. warmup schedules).

This change aligns schedule_free with the behavior of other stateful optimizers in Optax and ensures the first scheduled learning rate is applied as intended.