def context_managers(self, kwargs):
    mode = kwargs.pop("mode", None)
    if mode is None:
      return []
    elif mode == "eager":
      return [context.eager_mode()]
    elif mode == "graph":
      return [ops.Graph().as_default(), context.graph_mode()]
    else:
      raise ValueError(
          "Argument 'mode' must be either 'eager' or 'graph'. "
          f"Received: {mode}.")