# Path: tensorflow_constrained_optimization/python/train/constrained_optimizer.py
# class Formulation(tf.Module):
# class ConstrainedOptimizerV1(tf.compat.v1.train.Optimizer):
# class ConstrainedOptimizerV2(tf.keras.optimizers.Optimizer):
#   def __init__(self, name=None):
#   def state(self):
#   def create_state(self, num_constraints):
#   def is_state_created(self):
#   def get_loss_fn(self, minimization_problem):
#   def __init__(self,
#                formulation,
#                optimizer,
#                num_constraints=None,
#                constraint_optimizer=None,
#                name="ConstrainedOptimizerV1"):
#   def num_constraints(self):
#   def num_constraints(self, num_constraints):
#   def variables(self):
#   def trainable_variables(self):
#   def non_trainable_variables(self):
#   def compute_gradients(self,
#                         loss,
#                         var_list=None,
#                         gate_gradients=tf.compat.v1.train.Optimizer.GATE_OP,
#                         aggregation_method=None,
#                         colocate_gradients_with_ops=False,
#                         grad_loss=None):
#   def _create_slots(self, var_list):
#   def _prepare(self):
#   def _apply_dense(self, gradient, variable, *args, **kwargs):
#   def _apply_sparse(self, gradient, variable, *args, **kwargs):
#   def _resource_apply_dense(self, gradient, handle, *args, **kwargs):
#   def _resource_apply_sparse(self, gradient, handle, *args, **kwargs):
#   def __init__(self,
#                formulation,
#                optimizer,
#                num_constraints=None,
#                constraint_optimizer=None,
#                name="ConstrainedOptimizerV2"):
#   def num_constraints(self):
#   def num_constraints(self, num_constraints):
#   def variables(self):
#   def trainable_variables(self):
#   def non_trainable_variables(self):
#   def get_gradients(self, loss, params):
#   def _split_var_list(self, var_list):
#   def _compute_gradients(self, loss, var_list, grad_loss=None, tape=None):
#   def _create_slots(self, var_list):
#   def _prepare(self, var_list):
#   def _create_hypers(self):
#   def get_config(self):
#   def _resource_apply_dense(self, gradient, handle, *args, **kwargs):
#   def _resource_apply_sparse(self, gradient, handle, *args, **kwargs):
# 
# Path: tensorflow_constrained_optimization/python/train/lagrangian_optimizer.py
# def _project_multipliers_wrt_euclidean_norm(multipliers, radius):
#   def while_loop_condition(iteration, multipliers, inactive, old_inactive):
#   def while_loop_body(iteration, multipliers, inactive, old_inactive):
#   def __init__(self,
#                maximum_multiplier_radius=None,
#                dual_scale=1.0,
#                variable_fn=tf.Variable):
#   def state(self):
#   def _project_multipliers(self, multipliers):
#   def create_state(self, num_constraints):
#   def is_state_created(self):
#   def get_loss_fn(self, minimization_problem):
#       def loss_gradient_fn(objective, constraints, proxy_constraints,
#                            multipliers):
#         def gradient_fn(output_gradient):
#       def partial_loss_gradient_fn():
#       def loss_gradient_fn(objective):
# def create_lagrangian_loss(minimization_problem,
#                            maximum_multiplier_radius=None,
#                            dual_scale=1.0,
#                            variable_fn=tf.Variable):
#   def __init__(self,
#                optimizer,
#                num_constraints=None,
#                constraint_optimizer=None,
#                maximum_multiplier_radius=None,
#                name="LagrangianOptimizerV1"):
#   def __init__(self,
#                optimizer,
#                num_constraints=None,
#                constraint_optimizer=None,
#                maximum_multiplier_radius=None,
#                name="LagrangianOptimizerV2"):
# class _LagrangianFormulation(constrained_optimizer.Formulation):
# class LagrangianOptimizerV1(constrained_optimizer.ConstrainedOptimizerV1):
# class LagrangianOptimizerV2(constrained_optimizer.ConstrainedOptimizerV2):


import tensorflow as tf
from tensorflow_constrained_optimization.python.train import constrained_optimizer
from tensorflow_constrained_optimization.python.train import lagrangian_optimizer
      regularization_losses: This argument is ignored and can be left
        unspecified.

    Returns:
      `EstimatorSpec`.

    Raises:
      ValueError: If mode is not recognized or optimizer is not specified in
        TRAIN mode.
    """
    estimator_spec = self._create_no_op_estimator_spec(
        features=features,
        mode=mode,
        logits=logits,
        labels=labels)

    # When mode is PREDICT or EVAL, no modification needed to the base head.
    if (mode == tf.estimator.ModeKeys.PREDICT) or (
        mode == tf.estimator.ModeKeys.EVAL):
      return estimator_spec

    # When mode is TRAIN, replace train_op in estimator_spec.
    if mode == tf.estimator.ModeKeys.TRAIN:
      if optimizer is None:
        raise ValueError("You must provide an optimizer to the estimator.")
      # TODO: Add support for passing train_op_fn.

      # If the optimizer is not a `ConstrainedOptimizerV2` instance, then
      # create a `LagrangianOptimizer` that wraps the base heads's optimizer.
      if not isinstance(

# Path: tensorflow_constrained_optimization/python/train/constrained_optimizer.py
# class Formulation(tf.Module):
# class ConstrainedOptimizerV1(tf.compat.v1.train.Optimizer):
# class ConstrainedOptimizerV2(tf.keras.optimizers.Optimizer):
#   def __init__(self, name=None):
#   def state(self):
#   def create_state(self, num_constraints):
#   def is_state_created(self):
#   def get_loss_fn(self, minimization_problem):
#   def __init__(self,
#                formulation,
#                optimizer,
#                num_constraints=None,
#                constraint_optimizer=None,
#                name="ConstrainedOptimizerV1"):
#   def num_constraints(self):
#   def num_constraints(self, num_constraints):
#   def variables(self):
#   def trainable_variables(self):
#   def non_trainable_variables(self):
#   def compute_gradients(self,
#                         loss,
#                         var_list=None,
#                         gate_gradients=tf.compat.v1.train.Optimizer.GATE_OP,
#                         aggregation_method=None,
#                         colocate_gradients_with_ops=False,
#                         grad_loss=None):
#   def _create_slots(self, var_list):
#   def _prepare(self):
#   def _apply_dense(self, gradient, variable, *args, **kwargs):
#   def _apply_sparse(self, gradient, variable, *args, **kwargs):
#   def _resource_apply_dense(self, gradient, handle, *args, **kwargs):
#   def _resource_apply_sparse(self, gradient, handle, *args, **kwargs):
#   def __init__(self,
#                formulation,
#                optimizer,
#                num_constraints=None,
#                constraint_optimizer=None,
#                name="ConstrainedOptimizerV2"):
#   def num_constraints(self):
#   def num_constraints(self, num_constraints):
#   def variables(self):
#   def trainable_variables(self):
#   def non_trainable_variables(self):
#   def get_gradients(self, loss, params):
#   def _split_var_list(self, var_list):
#   def _compute_gradients(self, loss, var_list, grad_loss=None, tape=None):
#   def _create_slots(self, var_list):
#   def _prepare(self, var_list):
#   def _create_hypers(self):
#   def get_config(self):
#   def _resource_apply_dense(self, gradient, handle, *args, **kwargs):
#   def _resource_apply_sparse(self, gradient, handle, *args, **kwargs):
# 
# Path: tensorflow_constrained_optimization/python/train/lagrangian_optimizer.py
# def _project_multipliers_wrt_euclidean_norm(multipliers, radius):
#   def while_loop_condition(iteration, multipliers, inactive, old_inactive):
#   def while_loop_body(iteration, multipliers, inactive, old_inactive):
#   def __init__(self,
#                maximum_multiplier_radius=None,
#                dual_scale=1.0,
#                variable_fn=tf.Variable):
#   def state(self):
#   def _project_multipliers(self, multipliers):
#   def create_state(self, num_constraints):
#   def is_state_created(self):
#   def get_loss_fn(self, minimization_problem):
#       def loss_gradient_fn(objective, constraints, proxy_constraints,
#                            multipliers):
#         def gradient_fn(output_gradient):
#       def partial_loss_gradient_fn():
#       def loss_gradient_fn(objective):
# def create_lagrangian_loss(minimization_problem,
#                            maximum_multiplier_radius=None,
#                            dual_scale=1.0,
#                            variable_fn=tf.Variable):
#   def __init__(self,
#                optimizer,
#                num_constraints=None,
#                constraint_optimizer=None,
#                maximum_multiplier_radius=None,
#                name="LagrangianOptimizerV1"):
#   def __init__(self,
#                optimizer,
#                num_constraints=None,
#                constraint_optimizer=None,
#                maximum_multiplier_radius=None,
#                name="LagrangianOptimizerV2"):
# class _LagrangianFormulation(constrained_optimizer.Formulation):
# class LagrangianOptimizerV1(constrained_optimizer.ConstrainedOptimizerV1):
# class LagrangianOptimizerV2(constrained_optimizer.ConstrainedOptimizerV2):


# Path: tensorflow_constrained_optimization/python/rates/estimator_head.py
import tensorflow as tf
from tensorflow_constrained_optimization.python.train import constrained_optimizer
from tensorflow_constrained_optimization.python.train import lagrangian_optimizer

      regularization_losses: This argument is ignored and can be left
        unspecified.

    Returns:
      `EstimatorSpec`.

    Raises:
      ValueError: If mode is not recognized or optimizer is not specified in
        TRAIN mode.
    """
    estimator_spec = self._create_no_op_estimator_spec(
        features=features,
        mode=mode,
        logits=logits,
        labels=labels)

    # When mode is PREDICT or EVAL, no modification needed to the base head.
    if (mode == tf.estimator.ModeKeys.PREDICT) or (
        mode == tf.estimator.ModeKeys.EVAL):
      return estimator_spec

    # When mode is TRAIN, replace train_op in estimator_spec.
    if mode == tf.estimator.ModeKeys.TRAIN:
      if optimizer is None:
        raise ValueError("You must provide an optimizer to the estimator.")
      # TODO: Add support for passing train_op_fn.

      # If the optimizer is not a `ConstrainedOptimizerV2` instance, then
      # create a `LagrangianOptimizer` that wraps the base heads's optimizer.
      if not isinstance(

          optimizer, constrained_optimizer.ConstrainedOptimizerV2):