Adafactor accepts a list of dimensions not to factor or reduce. This will help Adafactor maintain identical behavior with models that stack variables together.
Adafactor accepts a list of dimensions not to factor or reduce.
This will help Adafactor maintain identical behavior with models that stack
variables together.

PiperOrigin-RevId: 364861973

tensorflow/mesh,tensorflow/mesh