File size: 24,949 Bytes
b100e1c
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
# Copyright 2022 The T5X Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Adafactor Optimizer.

Specialized Adafactor implementation for T5X with:
 - custom factorization specification rules.
 - support for stacked parameters from scanned layers and parameter fusions.

Why do we need custom factorization?  In the Adafactor paper, scalar, vector and
matrix parameters are considered. This is sufficiently general because higher
dimensional parameters can be reshaped. In practice, there are situations where
higher dimensional parameters are desirable.  For example, consider the
multi-headed attention. It has projection kernels.  This is naturally
represented as 3-dimensional array [d_model, num_head, head_dim]. Keeping the
3-dimensional structure can be beneficial for performance optimization, e.g., by
giving compilers additional degree of freedom to do layout optimization.

The default heuristic behavior for the second-moment estimator can lead to an
unexpected result because it assumes that the parameters are matrices (vectors
and scalars are not factored). The dimensions are sorted and the smaller
dimension is assigned to the row dim and the larger dim to the col dim (unless
the two largest dims have an equal size and then the original ordering of the
dimensions is used). Then `v_row` (i.e., the optimizer state for the row) is
obtained by removing the col dim. In other words, `rank(v_row) = rank(v) - 1`.
If the parameter is higher dimensional, v_row and v_col are higher dimensional.
Therefore, the outer product of v_row and v_col do not necessarily corresponds
to the row rank approximation that minimizes the generalized Kullback-Leibler
divergence (the original Adafactor formulation).

This Adafactor implementation generalized the default behavior such that we
obtain the correct second moment estimator even for higher dimensional
parameters.

"""
import enum
import re
from typing import Any, Mapping, Optional, Sequence, Tuple, Union

from absl import logging
from flax import struct
from flax.core import freeze
from flax.core import FrozenDict
from flax.core import unfreeze
from flax.serialization import from_state_dict
from flax.serialization import to_state_dict
from flax.traverse_util import flatten_dict
from flax.traverse_util import unflatten_dict
import jax
import jax.numpy as jnp
import numpy as np
from t5x import utils
from t5x.optimizers import OptimizerDef
from t5x.optimizers import OptimizerState

Dtype = Any


class FactorDim(enum.Enum):
  # Don't factorize this dimension.
  NONE = None
  # A batch-like dimension that we should not average over.
  BATCH = 1
  ROW = 2
  COLUMN = 3


# Sentinel value signifying the legacy heuristic factorization rule.
class HeuristicRule(enum.Enum):
  token = 1


HEURISTIC_RULE = HeuristicRule.token
FactorRule = Union[HeuristicRule, Tuple[FactorDim]]


def _restore(target, flat):
  state_dict = unflatten_dict({tuple(k.split('/')): v for k, v in flat.items()})
  if isinstance(target, FrozenDict):
    return freeze(state_dict)
  else:
    return state_dict


def _insert(tpl, idx, x):
  tmp = list(tpl)
  tmp.insert(idx, x)
  return tuple(tmp)


def standard_logical_factor_rules():
  return freeze({
      'vocab': FactorDim.COLUMN,
      'embed': FactorDim.ROW,
      'mlp': FactorDim.COLUMN,
      'heads': FactorDim.COLUMN,
      'kv': FactorDim.COLUMN,
      'joined_kv': FactorDim.COLUMN,
      'relpos_buckets': FactorDim.NONE,
      'layers': FactorDim.BATCH,  # used in scanned layers
      'stack': FactorDim.BATCH,  # used in stacked params
      # 'batch', 'length' should not occur in parameters
      'q_wi_fused': FactorDim.COLUMN,
      'o_wo_fused': FactorDim.COLUMN,
      'multiquery_heads': FactorDim.COLUMN,
      'kv_fused': FactorDim.COLUMN,
      'layer_norm_scale': FactorDim.NONE,
      'mlp_activations': FactorDim.COLUMN,
  })


def factor_name_to_factordim(name):
  if not isinstance(name, str):
    return name
  name = name.lower()
  return {
      'row': FactorDim.ROW,
      'col': FactorDim.COLUMN,
      'column': FactorDim.COLUMN,
      'batch': FactorDim.BATCH,
      'none': FactorDim.NONE,
      'unfactorized': FactorDim.NONE
  }[name]


class HParamMap:
  """Maps parameter path names to hparams.

  Names of parameters nested in a PyTree (e.g., an Optimizer) are formed by
  joining the names along the path to the parameter leaf with '/'.
  """

  def __init__(self, rules):
    self._rules = [(re.compile(r), p) for r, p in rules]

  def __getitem__(self, key: str) -> Any:
    for r, p in self._rules:
      if r.search(key):
        return p
    raise KeyError(f'No factor rule found for parameter: {key}')

  def __call__(self, params):
    """Returns a copy of the params with mapped hparams in leaves."""
    flat_state_dict = flatten_dict(to_state_dict(params))
    flat_rules_dict = {k: self['/'.join(k)] for k in flat_state_dict.keys()}
    return from_state_dict(params, unflatten_dict(flat_rules_dict))


@struct.dataclass
class _AdafactorHyperParams:
  """Hparams for Adafactor optimizer."""
  learning_rate: Optional[float]
  factored: bool
  multiply_by_parameter_scale: Union[bool, HParamMap]
  beta1: Optional[float]
  decay_rate: float
  step_offset: int
  clipping_threshold: Optional[float]
  weight_decay_rate: Optional[float]
  min_dim_size_to_factor: int
  epsilon1: float
  epsilon2: float
  factor_map: Optional[HParamMap] = None
  logical_factor_rules: Any = None
  weight_decay_rate_lr_exponent: Optional[float] = None
  global_norm_clip_threshold: Optional[float] = None
  max_parameter_scale: Optional[float] = None
  skip_nan_updates: Optional[bool] = False


@struct.dataclass
class _AdafactorParamState:
  v_row: np.ndarray  # used in normal factored version
  v_col: np.ndarray
  v: np.ndarray  # only used without factoring
  m: np.ndarray  # only used with momentum


class Adafactor(OptimizerDef):
  """Adafactor optimizer.

  Adafactor is described in https://arxiv.org/abs/1804.04235.
  """

  def __init__(self,
               learning_rate: Optional[float] = None,
               factored: bool = True,
               multiply_by_parameter_scale: Union[bool, HParamMap] = True,
               beta1: Optional[float] = None,
               decay_rate: float = 0.8,
               step_offset: int = 0,
               clipping_threshold: Optional[float] = 1.0,
               weight_decay_rate: Optional[float] = None,
               min_dim_size_to_factor: int = 128,
               epsilon1: float = 1e-30,
               epsilon2: float = 1e-3,
               dtype_momentum: Dtype = jnp.float32,
               factor_map: Optional[HParamMap] = None,
               logical_factor_rules: Optional[Mapping[str, FactorDim]] = None,
               weight_decay_rate_lr_exponent: Optional[float] = None,
               global_norm_clip_threshold: Optional[float] = None,
               max_parameter_scale: Optional[float] = None,
               skip_nan_updates: Optional[bool] = False):
    """Constructor for the Adafactor optimizer.


    Args:
      learning_rate: float: learning rate.  NB: the natural scale for adafactor
        LR is markedly different from Adam, one doesn't use the 1/sqrt(hidden)
        correction for this optimizer with attention-based models.
      factored: boolean: whether to use factored second-moment estimator for 2d
        variables.
      multiply_by_parameter_scale: boolean: if True, then scale provided
        learning_rate by parameter norm. if False, provided learning_rate is
        absolute step size.
      beta1: an optional float value between 0 and 1, enables momentum and uses
        extra memory if non-None! None by default.
      decay_rate: float: controls second-moment exponential decay schedule.
      step_offset: for finetuning, one may optionally set this to the starting
        step-number of the finetuning phase to reset the second moment
        accumulators after pretraining. Does not affect the momentum even if it
        was used during pretraining.
      clipping_threshold: an optional float >= 1, if None no update clipping.
      weight_decay_rate: optional rate at which to decay weights.
      min_dim_size_to_factor: only factor accumulator if two array dimensions
        are at least this size.
      epsilon1: Regularization constant for squared gradient.
      epsilon2: Regularization constant for parameter scale.
      dtype_momentum: dtype of momentum buffers.
      factor_map: hparam-map from key path to manual factorization rules.
      logical_factor_rules: factorization rules provided as a set of mappings
        from logical axis name to ROW, COLUMN, BATCH, or NONE.  Supercedes
        factor_map if `set_param_axes` is called.
      weight_decay_rate_lr_exponent: If present, weight decay rate is computed
        as (learning_rate ** weight_decay_rate_lr_exponent).  If
        weight_decay_rate is also present, then multiply by it.
      global_norm_clip_threshold: If set, will clip gradients by global norm
        before Adafactor stats are applied.
      max_parameter_scale: If set, clips the parameter scale to a maximum value,
        which helps prevent parameters from growing without bound.
      skip_nan_updates: If set, any parameter that would have been updated by a
        NaN value after a applying gradients will be kept with the earlier
        value it had.
    """
    if not factored and factor_map is not None:
      raise ValueError('Adafactor factored is False but factorization rules '
                       'have been provided.')
    if not isinstance(multiply_by_parameter_scale, (bool, HParamMap)):
      raise TypeError(
          '`multiply_by_parameter_scale` must be either bool or `HParamMap` '
          f'type. Got {type(multiply_by_parameter_scale)}')

    if not isinstance(factor_map, (type(None), HParamMap)):
      raise TypeError(
          '`factor_map` must be either None or `HParamMap` type. Got '
          f'{type(factor_map)}')

    hyper_params = _AdafactorHyperParams(
        learning_rate, factored, multiply_by_parameter_scale, beta1, decay_rate,
        step_offset, clipping_threshold, weight_decay_rate,
        min_dim_size_to_factor, epsilon1, epsilon2, factor_map,
        logical_factor_rules, weight_decay_rate_lr_exponent,
        global_norm_clip_threshold, max_parameter_scale, skip_nan_updates)
    self.dtype_momentum = jax.dtypes.canonicalize_dtype(dtype_momentum)
    super().__init__(hyper_params)

  @staticmethod
  def _decay_rate_pow(i: int, exponent: float = 0.8) -> float:
    """Default Adafactor second-moment decay schedule."""
    t = jnp.array(i, jnp.float32) + 1.0
    return 1.0 - t**(-exponent)

  @staticmethod
  def _parse_rule(
      rule: Optional[FactorRule],
      shape: Sequence[int],
      path: str,
      fallback_to_heuristics=True
  ) -> Tuple[Tuple[int, ...], Optional[Union[HeuristicRule, Tuple[Tuple[
      int, ...], Tuple[int, ...]]]]]:
    """Parses specification and return factored dims and dims for averaging.

    Adafactor needs to know the two largest dimensions to factorize along.
    Traditionally it used a heuristic, but we want finer control over these
    factorization dimensions.  Additionally, there are situations where
    parameters are batched together for e.g. scanned layers and QKV fusion,
    and we want to ensure that the scale updates and clipping thresholds are
    calculated _within_ each array and not across the entire batched array.

    Args:
      rule: the rule is either None (default to heuristic behavior) or a tuple
        of the same rank as the `param` array containing a FactorDim.ROW or
        FactorDim.COLUMN to mark dimensions to factorize in two row and column
        sets, and optionally dimensions marked FactorDim.BATCH to denote batched
        dimensions that should not be averaged over. e.g. (BATCH, ROW, COLUMN,
        COLUMN)
      shape: shape of the variable
      path: '/' joined parameter path.
      fallback_to_heuristics: whether to fallback to heuristic factorization
        rule. For most cases this should be set to `True`.

    Returns:
      tuple of: tuple of dimensions to average over, 2-tuple of dimensions to
      factorize over.
    """
    param_ndim = len(shape)

    if rule is None:
      # No factorization.
      return tuple(np.arange(param_ndim)), None

    if rule is HEURISTIC_RULE:
      if param_ndim > 2:
        raise ValueError(
            f'A parameter with rank strictly higher than 2 must have an '
            f'explicit factorization rule: {path}, {shape}')
      # Even if no explicit rule is provided for the param, we still want to
      # average over all the dimensions for computing the RMS scale.
      return tuple(np.arange(param_ndim)), HEURISTIC_RULE

    if len(rule) != param_ndim:
      raise ValueError(f'Factorization rule {rule} has incorrect rank '
                       f'for param of rank {param_ndim}: {path}, {shape}')

    row_dims = tuple(idx for idx, d in enumerate(rule) if d == FactorDim.ROW)
    col_dims = tuple(idx for idx, d in enumerate(rule) if d == FactorDim.COLUMN)
    batched_dims = tuple(
        idx for idx, d in enumerate(rule) if d == FactorDim.BATCH)
    averaging_dims = tuple(np.delete(np.arange(param_ndim), batched_dims))
    factor_dims = (row_dims, col_dims)
    if factor_dims == ((), ()):
      factor_dims = None

    if fallback_to_heuristics and param_ndim <= 2 and not batched_dims:
      logging.warning(
          'Since rank of parameter %s %d is less than or equal to 2, the '
          'factorization method falls back to heuristics and the provided '
          'factor rule %s is ignored.', path, param_ndim, rule)
      return tuple(np.arange(param_ndim)), HEURISTIC_RULE

    return averaging_dims, factor_dims

  def _factored_dims(
      self, shape: Sequence[int]) -> Optional[Tuple[Tuple[int], Tuple[int]]]:
    """Whether to use a factored second moment estimator.

    If there are not two dimensions of size >= min_dim_size_to_factor, then we
    do not factor. If we do factor the accumulator, then this function returns a
    tuple of the two largest axes to reduce over.

    Args:
      shape: a Shape

    Returns:
      None or a tuple of ints
    """
    if not self.hyper_params.factored or len(shape) < 2:
      return None
    sorted_dims = np.argsort(shape)
    if shape[sorted_dims[-2]] < self.hyper_params.min_dim_size_to_factor:
      return None
    return (int(sorted_dims[-2]),), (int(sorted_dims[-1]),)

  def init_param_state(self, param, path):
    shape = param.shape
    state = {k: jnp.zeros((1,)) for k in ['v_row', 'v_col', 'v', 'm']}
    if self.hyper_params.factored:
      factor_rule = (
          self.hyper_params.factor_map[path]
          if self.hyper_params.factor_map else HEURISTIC_RULE)
    else:
      factor_rule = None
    _, factored_dims = self._parse_rule(factor_rule, param.shape, path)
    if factored_dims is HEURISTIC_RULE:
      factored_dims = self._factored_dims(shape)
    if factored_dims is not None:
      d1, d0 = factored_dims
      vr_shape = np.delete(shape, d0)
      vc_shape = np.delete(shape, d1)
      state['v_row'] = jnp.zeros(vr_shape, dtype=jnp.float32)
      state['v_col'] = jnp.zeros(vc_shape, dtype=jnp.float32)
    else:
      state['v'] = jnp.zeros(param.shape, dtype=jnp.float32)
    if self.hyper_params.beta1 is not None:
      state['m'] = jnp.zeros(param.shape, dtype=self.dtype_momentum)
    return _AdafactorParamState(**state)

  def init_state(self, params):
    params_flat = utils.flatten_dict_string_keys(params)
    param_states_flat = [
        self.init_param_state(param, path)
        for path, param in params_flat.items()
    ]
    param_states_flat = {
        k: v for k, v in zip(params_flat.keys(), param_states_flat)
    }
    param_states = _restore(params, param_states_flat)
    state = OptimizerState(jnp.asarray(0, dtype=jnp.int32), param_states)
    return state

  def apply_param_gradient(self, step, hyper_params, param, state, grad, path):
    assert hyper_params.learning_rate is not None, 'no learning rate provided.'
    learning_rate = hyper_params.learning_rate
    beta1 = hyper_params.beta1
    decay_rate = hyper_params.decay_rate
    step_offset = hyper_params.step_offset
    multiply_by_parameter_scale = hyper_params.multiply_by_parameter_scale
    max_parameter_scale = hyper_params.max_parameter_scale
    clipping_threshold = hyper_params.clipping_threshold
    weight_decay_rate = hyper_params.weight_decay_rate
    epsilon1 = hyper_params.epsilon1
    epsilon2 = hyper_params.epsilon2
    if hyper_params.weight_decay_rate_lr_exponent:
      weight_decay_rate = (
          (weight_decay_rate or 1.0) *
          learning_rate**hyper_params.weight_decay_rate_lr_exponent)

    if self.hyper_params.factored:
      factor_rule = (
          self.hyper_params.factor_map[path]
          if self.hyper_params.factor_map else HEURISTIC_RULE)
    else:
      factor_rule = None
    averaging_dims, factored_dims = self._parse_rule(factor_rule, param.shape,
                                                     path)

    grad = grad.astype(jnp.float32)

    updates = {k: jnp.zeros((1,)) for k in ['v_row', 'v_col', 'v', 'm']}
    decay_rate = self._decay_rate_pow(step - step_offset, exponent=decay_rate)
    update_scale = learning_rate

    if isinstance(multiply_by_parameter_scale, HParamMap):
      multiply_by_parameter_scale = multiply_by_parameter_scale[path]
    if multiply_by_parameter_scale:
      param_scale = jnp.sqrt(
          jnp.mean(param * param, axis=averaging_dims, keepdims=True))
      # Clip param_scale to a minimum value of epsilon2.
      param_scale = jnp.maximum(param_scale, epsilon2)
      # Clip param_scale to a maximum value, if specified.
      if max_parameter_scale is not None:
        param_scale = jnp.minimum(param_scale, max_parameter_scale)
      update_scale *= param_scale
    mixing_rate = 1.0 - decay_rate

    grad_sqr = grad * grad + epsilon1
    if factored_dims is HEURISTIC_RULE:
      factored_dims = self._factored_dims(param.shape)
    if factored_dims is not None:
      d1, d0 = factored_dims
      new_v_row = (
          decay_rate * state.v_row + mixing_rate * jnp.mean(grad_sqr, axis=d0))
      new_v_col = (
          decay_rate * state.v_col + mixing_rate * jnp.mean(grad_sqr, axis=d1))
      updates['v_row'] = new_v_row
      updates['v_col'] = new_v_col
      reduced_d1 = tuple(d - len([e for e in d0 if e < d]) for d in d1)

      row_col_mean = jnp.mean(new_v_row, axis=reduced_d1, keepdims=True)
      row_factor = (new_v_row / row_col_mean)**-0.5
      col_factor = (new_v_col)**-0.5
      y = (
          grad * jnp.expand_dims(row_factor, axis=d0) *
          jnp.expand_dims(col_factor, axis=d1))
    else:
      new_v = decay_rate * state.v + mixing_rate * grad_sqr
      updates['v'] = new_v
      y = grad * (new_v)**-0.5

    if clipping_threshold is not None:
      clipping_denom = (
          jnp.maximum(
              1.0,
              jnp.sqrt(jnp.mean(y * y, axis=averaging_dims, keepdims=True)) /
              clipping_threshold))
      y /= clipping_denom

    subtrahend = update_scale * y
    if beta1 is not None:
      new_m = beta1 * state.m + (1.0 - beta1) * subtrahend
      subtrahend = new_m
      updates['m'] = new_m.astype(self.dtype_momentum)

    if weight_decay_rate is not None:
      new_param = (1.0 - weight_decay_rate) * param - subtrahend
    else:
      new_param = param - subtrahend

    if hyper_params.skip_nan_updates:
      updates['v_row'] = jnp.where(
          jnp.isnan(updates['v_row']), state.v_row, updates['v_row'])
      updates['v_col'] = jnp.where(
          jnp.isnan(updates['v_col']), state.v_col, updates['v_col'])
      updates['v'] = jnp.where(jnp.isnan(updates['v']), state.v, updates['v'])
      updates['m'] = jnp.where(jnp.isnan(updates['m']), state.m, updates['m'])
      new_param = jnp.where(jnp.isnan(new_param), param, new_param)
    new_state = _AdafactorParamState(**updates)

    return new_param.astype(param.dtype), new_state

  def apply_gradient(self, hyper_params, params, state, grads):
    """Applies a gradient for a set of parameters.

    Args:
      hyper_params: a named tuple of hyper parameters.
      params: the parameters that should be updated.
      state: a named tuple containing the state of the optimizer
      grads: the gradient tensors for the parameters.

    Returns:
      A tuple containing the new parameters and the new optimizer state.
    """
    step = state.step
    # We assume that params, param_states, and grads are all dict-like here.
    params_flat_dict = utils.flatten_dict_string_keys(params)
    params_paths = params_flat_dict.keys()
    params_flat = params_flat_dict.values()
    # extra paranoia to guarantee identical value ordering
    states_flat = utils.flatten_dict_string_keys(state.param_states)
    states_flat = [states_flat[k] for k in params_paths]
    grads_flat = utils.flatten_dict_string_keys(grads)
    grads_flat = [grads_flat[k] for k in params_paths]

    if hyper_params.global_norm_clip_threshold:
      # Paper: http://proceedings.mlr.press/v28/pascanu13.pdf
      # TF: https://www.tensorflow.org/api_docs/python/tf/clip_by_global_norm
      squared_l2_norms = [jnp.sum(jnp.square(g)) for g in grads_flat]
      global_norm = jnp.sqrt(jnp.sum(jnp.array(squared_l2_norms)))
      scale = hyper_params.global_norm_clip_threshold * jnp.minimum(
          1.0 / hyper_params.global_norm_clip_threshold, 1.0 / global_norm)
      grads_flat = [g * scale for g in grads_flat]

    out = [
        self.apply_param_gradient(step, hyper_params, param, state, grad, path)
        for param, state, grad, path in zip(params_flat, states_flat,
                                            grads_flat, params_paths)
    ]

    new_params_flat, new_states_flat = list(zip(*out)) if out else ((), ())
    new_params_flat = {k: v for k, v in zip(params_paths, new_params_flat)}
    new_states_flat = {k: v for k, v in zip(params_paths, new_states_flat)}
    new_params = _restore(params, new_params_flat)
    new_param_states = _restore(params, new_states_flat)
    new_state = OptimizerState(step + 1, new_param_states)

    return new_params, new_state

  def set_param_axes(self, param_logical_axes):
    """Sets Adafactor factorization map from logical axis names tree."""
    logical_factor_rules = self.hyper_params.logical_factor_rules
    if logical_factor_rules is None:
      return

    # pylint:disable=invalid-name
    NONE = FactorDim.NONE
    COLUMN = FactorDim.COLUMN
    ROW = FactorDim.ROW

    # pylint:enable=invalid-name

    def apply_rules(axes):
      # Partially factorized params are marked as unfactorized, preserving
      # only BATCH axis annotations. We also check for incompletely factorized
      # params that have ROW, COLUMN but also accidental NONE dimensions and
      # raise an error in that case.
      axis_rules = tuple(logical_factor_rules[x] for x in axes)
      axis_rules = tuple(factor_name_to_factordim(x) for x in axis_rules)
      if ROW in axis_rules and COLUMN in axis_rules and NONE in axis_rules:
        raise ValueError(f'Incomplete adafactor spec {axis_rules} for {axes}!')
      if ROW not in axis_rules or COLUMN not in axis_rules:
        axis_rules = tuple(
            NONE if x in (ROW, COLUMN) else x for x in axis_rules)
      return axis_rules

    factor_map = jax.tree_map(apply_rules, param_logical_axes)
    factor_map = utils.flatten_dict_string_keys(factor_map)

    self.hyper_params = self.hyper_params.replace(factor_map=factor_map)

  def derive_logical_axes(self, optimizer_state, param_logical_axes):
    """Derives optimizer logical partitioning from model logical partitions."""
    optimizer_logical_axes = jax.tree_map(lambda x: None,
                                          optimizer_state.state_dict())
    optimizer_logical_axes['target'] = param_logical_axes

    def factor_rule(logical_axes, adafactor_leaf):
      return dict(
          v_row=None,
          v_col=None,
          v=logical_axes if adafactor_leaf['v'].shape != (1,) else None,
          m=logical_axes if self.hyper_params.beta1 else None)

    optimizer_logical_axes['state']['param_states'] = jax.tree_map(
        factor_rule, unfreeze(param_logical_axes),
        optimizer_state.state_dict()['state']['param_states'])

    return optimizer_state.restore_state(unfreeze(optimizer_logical_axes))