class_docstr: Base context class that contains runtime information for any Determined
  workflow that uses the ``tf.keras`` API.
class_docstr_tok: ["Base","context","class","that","contains","runtime","information","for","any","Determined","workflow","that","uses","the","``","tf.keras","``","API","."]
docstr: This should be user to wrap ``tf.keras.optimizers.Optimizer`` objects. Users
  should use the output use the output of this wrapper as the new instance of
  their optimizer. If users create multiple optimizers, users should wrap each
  optimizer independently.
  
  Args:
      optimizer: tf.keras.optimizers.Optimizer
docstr_tok: ["This","should","be","user","to","wrap","``","tf.keras.optimizers.Optimizer","``","objects",".","Users","should","use","the","output","use","the","output","of","this","wrapper","as","the","new","instance","of","their","optimizer",".","If","users","create","multiple","optimizers",",","users","should","wrap","each","optimizer","independently",".","Args",":","optimizer",":","tf.keras.optimizers.Optimizer"]
code: def wrap_optimizer(
      self, optimizer: tf.keras.optimizers.Optimizer
  ) -> tf.keras.optimizers.Optimizer:
      """
      This should be user to wrap ``tf.keras.optimizers.Optimizer`` objects. Users
      should use the output use the output of this wrapper as the new instance of
      their optimizer. If users create multiple optimizers, users should wrap each
      optimizer independently.
  
      Args:
          optimizer: tf.keras.optimizers.Optimizer
      """
      if not self.env.managed_training:
          return optimizer
  
      logger.debug(f"Processing wrapped optimizer {optimizer}.")
      if self.distributed.size == 1:
          self._wrapped_optimizers.append(optimizer)
          return optimizer
  
      hvd.require_horovod_type(
          "tensorflow.keras",
          "TFKerasTrialContext.wrap_optimizer was called.",
      )
      if optimizer == self._compiled_optimizer:
          logger.debug(
              "Skipping wrapping optimizer as it was already wrapped during the compile call."
          )
          wrapped_optimizer = optimizer
      else:
          wrapped_optimizer = (
              self._get_horovod_optimizer_if_using_horovod(
                  optimizer=optimizer,
              )
          )
      self._wrapped_optimizers.append(wrapped_optimizer)
  
      return wrapped_optimizer
  
imports: import inspect
  import logging
  from typing import Any, Callable, Dict, List, NamedTuple, Optional, Union, cast
  import tensorflow
  import determined
  from determined import errors, keras, util
  from determined.common import check
  from determined.horovod import hvd
label_desc: Use image node_id 8 for calling the TFKerasTrialContext obj's underlying member method code with example usage: obj.wrap_optimizer(optimizer) and returns: wrapped_optimizer, optimizer, optimizer