tfm.optimization.register_optimizer_cls

Register customize optimizer cls.

tfm.optimization.register_optimizer_cls(
    key: str,
    optimizer_config_cls: Union[tf.keras.optimizers.Optimizer, tf.keras.optimizers.legacy.Optimizer,
        tf.keras.optimizers.experimental.Optimizer],
    use_legacy_optimizer: bool = True
)

The user will still need to subclass data classes in configs.optimization_config to be used with OptimizerFactory.

key A string to that the optimizer_config_cls is registered with.
optimizer_config_cls A class which inherits tf.keras.optimizers.Optimizer.
use_legacy_optimizer A boolean that indicates if using legacy optimizers.