flax

7a0e1df8 - Add offset argument to Adafactor optimizer for fine-tuning schedules.

Add offset argument to Adafactor optimizer for fine-tuning schedules. This argument can adjust the Adafactor mixing-rate schedule, which can be useful during finetuning of a pretrained model. PiperOrigin-RevId: 341521827