Class PartitionedGradientTransformation

axlearn/common/optimizer_base.py:71–82

An optax-style optimizer with a function to partition the inputs across devices. For new optimizers, using `UpdateTransformation` is preferred instead because it supports more types of optimizers and allows better reuse of functionality across different optimizers. Despite this, there

class PartitionedGradientTransformation(NamedTuple):
    """An optax-style optimizer with a function to partition the inputs across devices.

    For new optimizers, using `UpdateTransformation` is preferred instead because it supports
    more types of optimizers and allows better reuse of functionality across different optimizers.

    Despite this, there are no plans to stop supporting this class.
    """

    init: TransformInitFn
    update: TransformUpdateFn
    partition: TransformPartitionSpecFn