# Path: meta_dataset/learners/base.py
# class Learner(object):
# class EpisodicLearner(Learner):
# class BatchLearner(Learner):
#   def __init__(
#       self,
#       is_training,
#       logit_dim,
#       transductive_batch_norm,
#       backprop_through_moments,
#       embedding_fn,
#       input_shape,
#   ):
#   def build(self):
#   def compute_regularizer(self):
#   def compute_loss(self, onehot_labels, predictions):
#   def compute_accuracy(self, onehot_labels, predictions):
#   def forward_pass(self, data):


import gin.tf
import tensorflow as tf
from meta_dataset.learners import base as learner_base
# Lint as: python2, python3
"""Abstract experimental learners that use `ReparameterizableModule`s."""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function




class NotBuiltError(RuntimeError):

  def __init__(self):
    super(NotBuiltError, self).__init__(
        'The `build` method of `ExperimentalLearner` must be called before '
        'accessing its variables.')


def class_specific_data(onehot_labels, data, num_classes, axis=0):
  # TODO(eringrant): Deal with case of no data for a class in [1...num_classes].
  data_shape = [s for i, s in enumerate(data.shape) if i != axis]
  labels = tf.argmax(onehot_labels, axis=-1)
  class_idx = [tf.where(tf.equal(labels, i)) for i in range(num_classes)]
  return [
      tf.reshape(tf.gather(data, idx, axis=axis), [-1] + data_shape)
      for idx in class_idx
  ]


@gin.configurable

# Path: meta_dataset/learners/base.py
# class Learner(object):
# class EpisodicLearner(Learner):
# class BatchLearner(Learner):
#   def __init__(
#       self,
#       is_training,
#       logit_dim,
#       transductive_batch_norm,
#       backprop_through_moments,
#       embedding_fn,
#       input_shape,
#   ):
#   def build(self):
#   def compute_regularizer(self):
#   def compute_loss(self, onehot_labels, predictions):
#   def compute_accuracy(self, onehot_labels, predictions):
#   def forward_pass(self, data):


# Path: meta_dataset/learners/experimental/base.py
import gin.tf
import tensorflow as tf
from meta_dataset.learners import base as learner_base

# Lint as: python2, python3
"""Abstract experimental learners that use `ReparameterizableModule`s."""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function




class NotBuiltError(RuntimeError):

  def __init__(self):
    super(NotBuiltError, self).__init__(
        'The `build` method of `ExperimentalLearner` must be called before '
        'accessing its variables.')


def class_specific_data(onehot_labels, data, num_classes, axis=0):
  # TODO(eringrant): Deal with case of no data for a class in [1...num_classes].
  data_shape = [s for i, s in enumerate(data.shape) if i != axis]
  labels = tf.argmax(onehot_labels, axis=-1)
  class_idx = [tf.where(tf.equal(labels, i)) for i in range(num_classes)]
  return [
      tf.reshape(tf.gather(data, idx, axis=axis), [-1] + data_shape)
      for idx in class_idx
  ]


@gin.configurable

class ExperimentalLearner(learner_base.Learner):