# Path: meta_dataset/learners/base_test.py
# IMAGE_HEIGHT = IMAGE_WIDTH = 84
# IMAGE_SHAPE = (IMAGE_HEIGHT, IMAGE_WIDTH, 3)
# SUPPORT_SIZE, QUERY_SIZE = 5, 15
# NUM_CLASSES = 5
# VALID_ABSTRACT_LEARNER_INIT_ARGS = {
#     'is_training': True,
#     'logit_dim': NUM_CLASSES,
#     'transductive_batch_norm': True,
#     'backprop_through_moments': True,
#     'input_shape': IMAGE_SHAPE,
# }
# VALID_LEARNER_INIT_ARGS = {
#     **VALID_ABSTRACT_LEARNER_INIT_ARGS, 'embedding_fn':
#         functools.partial(
#             functional_backbones.four_layer_convnet, weight_decay=0.01)
# }
# class MockEmbedding(reparameterizable_backbones.ConvNet):
# class MockEpisode(
#     collections.namedtuple(
#         'MockEpisode', 'support_images, query_images, '
#         'support_labels, query_labels')):
# class MockBatch(collections.namedtuple('MockBatch', 'images, labels')):
# class TestLearner(tf.test.TestCase):
# class TestBatchLearner(TestLearner):
# class TestEpisodicLearner(TestLearner):
#   def __init__(self, keep_spatial_dims=False):
#   def way(self):
#   def onehot_support_labels(self):
#   def onehot_query_labels(self):
#   def labels(self):
#   def onehot_labels(self):
#   def way(self):
#   def onehot_labels(self):
#   def learner_cls(self):
#   def learner_kwargs(self):
#   def set_up_learner(self):
#   def testForwardPass(self):
#   def testComputeLoss(self):
#   def testComputeAccuracy(self):
#   def testLearnerInitRandomAccuracy(self):
#   def testLearnerImprovement(self):
#   def testLearnerConvergence(self):
#   def random_data(self):
#   def random_data(self):
# 
# Path: meta_dataset/learners/optimization_learners.py
# def get_embeddings_vars_copy_ops(embedding_vars_dict, make_copies):
# def get_fc_vars_copy_ops(fc_weights, fc_bias, make_copies):
# def gradient_descent_step(loss,
#                           variables,
#                           stop_grads,
#                           allow_grads_to_batch_norm_vars,
#                           learning_rate,
#                           get_update_ops=True):
#   def _apply_grads(variables, grads):
#   def __init__(self,
#                num_finetune_steps,
#                finetune_lr,
#                debug_log=False,
#                finetune_all_layers=False,
#                finetune_with_adam=False,
#                **kwargs):
#   def compute_logits(self, data):
#   def _get_finetune_op(self,
#                        data,
#                        embedding_vars_keys,
#                        embedding_vars,
#                        vars_to_finetune,
#                        support_embeddings=None):
#   def _fc_layer(self, embedding):
#   def __init__(self,
#                num_update_steps,
#                additional_evaluation_update_steps,
#                first_order,
#                alpha,
#                adapt_batch_norm,
#                zero_fc_layer,
#                proto_maml_fc_layer_init,
#                classifier_weight_decay,
#                debug=False,
#                **kwargs):
#   def proto_maml_fc_weights(self, prototypes, zero_pad_to_max_way=False):
#   def proto_maml_fc_bias(self, prototypes, zero_pad_to_max_way=False):
#   def forward_pass(self, data):
#     def _cond(step, *args):
#     def _body(step, *args):
#   def __init__(self, num_steps, lr, film_init, debug_log=False, **kwargs):
#   def forward_pass(self, data):
#   def _forward_pass(self, images, params=None, moments=None):
#   def _compute_prototype_loss(self,
#                               embeddings,
#                               labels,
#                               labels_one_hot,
#                               prototypes=None):
#   def _compute_losses(self, data, compute_on_query=False):
#   def _get_train_op(self, data):
# class BaselineFinetuneLearner(baseline_learners.BaselineLearner):
# class OptimizationLearner(learner_base.EpisodicLearner):
# class MAMLLearner(OptimizationLearner):
# class FLUTEFiLMLearner(learner_base.EpisodicLearner):


import gin.tf
import tensorflow.compat.v1 as tf
from meta_dataset.learners import base_test
from meta_dataset.learners import optimization_learners
tf.disable_eager_execution()

BASELINE_FINETUNE_ARGS = {
    **base_test.VALID_LEARNER_INIT_ARGS,
    'knn_in_fc': False,
    'knn_distance': False,
    'cosine_classifier': False,
    'cosine_logits_multiplier': 1.0,
    'use_weight_norm': False,
    'is_training': False,
}

gin.bind_parameter('MAMLLearner.classifier_weight_decay', 0.01)


class BaselineFinetuneTest():

  def testLearnerConvergence(self):
    # `BaselineFinetuneLearner` differs from `BaselineLearner` only at
    # evaluation time.
    pass

  def testLearnerImprovement(self):
    # `BaselineFinetuneLearner` differs from `BaselineLearner` only at
    # evaluation time.
    pass


class BaselineFinetuneAllLayersAdamTest(BaselineFinetuneTest,
                                        base_test.TestEpisodicLearner):

# Path: meta_dataset/learners/base_test.py
# IMAGE_HEIGHT = IMAGE_WIDTH = 84
# IMAGE_SHAPE = (IMAGE_HEIGHT, IMAGE_WIDTH, 3)
# SUPPORT_SIZE, QUERY_SIZE = 5, 15
# NUM_CLASSES = 5
# VALID_ABSTRACT_LEARNER_INIT_ARGS = {
#     'is_training': True,
#     'logit_dim': NUM_CLASSES,
#     'transductive_batch_norm': True,
#     'backprop_through_moments': True,
#     'input_shape': IMAGE_SHAPE,
# }
# VALID_LEARNER_INIT_ARGS = {
#     **VALID_ABSTRACT_LEARNER_INIT_ARGS, 'embedding_fn':
#         functools.partial(
#             functional_backbones.four_layer_convnet, weight_decay=0.01)
# }
# class MockEmbedding(reparameterizable_backbones.ConvNet):
# class MockEpisode(
#     collections.namedtuple(
#         'MockEpisode', 'support_images, query_images, '
#         'support_labels, query_labels')):
# class MockBatch(collections.namedtuple('MockBatch', 'images, labels')):
# class TestLearner(tf.test.TestCase):
# class TestBatchLearner(TestLearner):
# class TestEpisodicLearner(TestLearner):
#   def __init__(self, keep_spatial_dims=False):
#   def way(self):
#   def onehot_support_labels(self):
#   def onehot_query_labels(self):
#   def labels(self):
#   def onehot_labels(self):
#   def way(self):
#   def onehot_labels(self):
#   def learner_cls(self):
#   def learner_kwargs(self):
#   def set_up_learner(self):
#   def testForwardPass(self):
#   def testComputeLoss(self):
#   def testComputeAccuracy(self):
#   def testLearnerInitRandomAccuracy(self):
#   def testLearnerImprovement(self):
#   def testLearnerConvergence(self):
#   def random_data(self):
#   def random_data(self):
# 
# Path: meta_dataset/learners/optimization_learners.py
# def get_embeddings_vars_copy_ops(embedding_vars_dict, make_copies):
# def get_fc_vars_copy_ops(fc_weights, fc_bias, make_copies):
# def gradient_descent_step(loss,
#                           variables,
#                           stop_grads,
#                           allow_grads_to_batch_norm_vars,
#                           learning_rate,
#                           get_update_ops=True):
#   def _apply_grads(variables, grads):
#   def __init__(self,
#                num_finetune_steps,
#                finetune_lr,
#                debug_log=False,
#                finetune_all_layers=False,
#                finetune_with_adam=False,
#                **kwargs):
#   def compute_logits(self, data):
#   def _get_finetune_op(self,
#                        data,
#                        embedding_vars_keys,
#                        embedding_vars,
#                        vars_to_finetune,
#                        support_embeddings=None):
#   def _fc_layer(self, embedding):
#   def __init__(self,
#                num_update_steps,
#                additional_evaluation_update_steps,
#                first_order,
#                alpha,
#                adapt_batch_norm,
#                zero_fc_layer,
#                proto_maml_fc_layer_init,
#                classifier_weight_decay,
#                debug=False,
#                **kwargs):
#   def proto_maml_fc_weights(self, prototypes, zero_pad_to_max_way=False):
#   def proto_maml_fc_bias(self, prototypes, zero_pad_to_max_way=False):
#   def forward_pass(self, data):
#     def _cond(step, *args):
#     def _body(step, *args):
#   def __init__(self, num_steps, lr, film_init, debug_log=False, **kwargs):
#   def forward_pass(self, data):
#   def _forward_pass(self, images, params=None, moments=None):
#   def _compute_prototype_loss(self,
#                               embeddings,
#                               labels,
#                               labels_one_hot,
#                               prototypes=None):
#   def _compute_losses(self, data, compute_on_query=False):
#   def _get_train_op(self, data):
# class BaselineFinetuneLearner(baseline_learners.BaselineLearner):
# class OptimizationLearner(learner_base.EpisodicLearner):
# class MAMLLearner(OptimizationLearner):
# class FLUTEFiLMLearner(learner_base.EpisodicLearner):


# Path: meta_dataset/learners/optimization_learners_test.py
import gin.tf
import tensorflow.compat.v1 as tf
from meta_dataset.learners import base_test
from meta_dataset.learners import optimization_learners

tf.disable_eager_execution()

BASELINE_FINETUNE_ARGS = {
    **base_test.VALID_LEARNER_INIT_ARGS,
    'knn_in_fc': False,
    'knn_distance': False,
    'cosine_classifier': False,
    'cosine_logits_multiplier': 1.0,
    'use_weight_norm': False,
    'is_training': False,
}

gin.bind_parameter('MAMLLearner.classifier_weight_decay', 0.01)


class BaselineFinetuneTest():

  def testLearnerConvergence(self):
    # `BaselineFinetuneLearner` differs from `BaselineLearner` only at
    # evaluation time.
    pass

  def testLearnerImprovement(self):
    # `BaselineFinetuneLearner` differs from `BaselineLearner` only at
    # evaluation time.
    pass


class BaselineFinetuneAllLayersAdamTest(BaselineFinetuneTest,
                                        base_test.TestEpisodicLearner):

  learner_cls = optimization_learners.BaselineFinetuneLearner