Training arguments for flax

Noticed there is no gradient accumulation option for the flax text classification example.

The language modeling example uses TrainingArguments, which includes gradient accumulation and other options. Is there any plan to update the text classification example?