Compose a rule with trailing transforms, and scale by a schedule

[jessegrabowski]

chain already worked with a rule at the head — the body threads whatever it's handed into the first callable — but the signature said transforms only, and mypy rejected chain(adam(1e-3), clip_by_global_norm(1.0), scale(0.5)). That's overloads now. head stays unannotated: any type wide enough to accept both is a union, and a union return won't pass to compile_train, so that one head(...) call is unchecked.

scale_by_schedule schedules the rate after the rule instead of inside it. Those are different graphs once a clip sits between the two — the clip bounds a step at unit rate, so its bound stays in gradient units rather than moving with whatever the schedule currently holds. It owns its clock, and compile_train advances it.

Closes #81
Closes #82

[codecov-commenter]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 97.52%. Comparing base (497d21b) to head (bcd6e04).

Additional details and impacted files

@@           Coverage Diff           @@
##             main     #108   +/-   ##
=======================================
  Coverage   97.52%   97.52%           
=======================================
  Files          61       61           
  Lines        3069     3076    +7     
=======================================
+ Hits         2993     3000    +7     
  Misses         76       76           

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

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.