Adafactor HParamMap factor rule definitions missing?

Hi,

I am trying to run the sample code for `scalable_t5` on a TPU VM (modified to a span corruption task). However, I am running into the following error message:

```
ValueError: A parameter with rank strictly higher than 2 must have an explicit factorization rule: decoder/layers/encode
r_decoder_attention/key/kernel, (768, 12, 12, 64)
```

From what I can tell, it looks like some additional configuration should be provided in the Gin files for configuring Adafactor?

Thanks!