

TITLE
[WIP][examples/flax] add gradient accumulation

COMMENTS
4

REACTIONS
+1: 0
-1: 0
laugh: 0
hooray: 0
heart: 0
rocket: 0
eyes: 0

BODY
# What does this PR do?

Adds gradient accumulation in flax language modeling scripts.