# Path: meta_dataset/learners/experimental/base.py
# class NotBuiltError(RuntimeError):
# class ExperimentalLearner(learner_base.Learner):
# class ExperimentalEpisodicLearner(ExperimentalLearner,
#                                   learner_base.EpisodicLearner):
# class ExperimentalBatchLearner(ExperimentalLearner, learner_base.BatchLearner):
#   def __init__(self):
# def class_specific_data(onehot_labels, data, num_classes, axis=0):
#   def __init__(self, **kwargs):
#   def compute_regularizer(self, onehot_labels, predictions):
#   def build(self):
#   def variables(self):
#   def trainable_variables(self):
# 
# Path: meta_dataset/models/experimental/reparameterizable_backbones.py
# class ReparameterizableBackbone(tf.keras.Sequential,
#                                 reparameterizable_base.ReparameterizableModule):
# class ReparameterizableSpatialBackbone(ReparameterizableBackbone):
# class LinearModel(ReparameterizableBackbone):
# class FullyConnectedNet(ReparameterizableBackbone):
# class ConvNet(ReparameterizableSpatialBackbone):
# class RelationNetConvNet(ReparameterizableSpatialBackbone):
# class RelationModule(ReparameterizableBackbone):
# class BasicBlock(tf.keras.Model):
# class BasicBlockV1(BasicBlock):
# class BasicBlockV2(BasicBlock):
# class BottleneckBlock(tf.keras.Model):
# class BottleneckBlockV1(BottleneckBlock):
# class BottleneckBlockV2(BottleneckBlock):
# class ResNet(ReparameterizableSpatialBackbone):
# class ResNet18(ResNet):
# class ResNet18V1(ResNet18):
# class ResNet18V2(ResNet18):
# class ResNet34(ResNet):
# class ResNet34V1(ResNet34):
# class ResNet34V2(ResNet34):
# class ResNet50(ResNet):
# class ResNet50V1(ResNet50):
# class ResNet50V2(ResNet50):
# class ResNet101(ResNet):
# class ResNet101V1(ResNet101):
# class ResNet101V2(ResNet101):
# class ResNet152(ResNet):
# class ResNet152V1(ResNet152):
# class ResNet152V2(ResNet152):
# class GNResNet(ResNet):
# class GNBasicBlockV2(BasicBlockV2):
# class GNResNet18V2(GNResNet):
# class GNBottleneckBlockV2(BottleneckBlockV2):
# class GNResNet50V2(GNResNet):
# class WideResNet(ResNet):
# class WideResNet16(WideResNet):
# class WideResNet16V1(WideResNet16):
# class WideResNet16V2(WideResNet16):
# class WideResNet28(WideResNet):
# class WideResNet28V1(WideResNet28):
# class WideResNet28V2(WideResNet28):
#   def __init__(self,
#                output_dim=None,
#                kernel_initializer=tf.keras.initializers.he_normal(),
#                kernel_regularizer=tf.keras.regularizers.l2(),
#                parameter_adapter=None,
#                name=None):
#   def build(self, input_shape):
#   def stack(self):
#   def output_layers(self):
#   def __init__(self, keep_spatial_dims=False, **kwargs):
#   def output_layers(self):
#   def __init__(self, output_dim, **kwargs):
#   def stack(self):
#   def __init__(self, num_hiddens_per_layer, **kwargs):
#   def stack(self):
#   def __init__(self, num_filters_per_layer, **kwargs):
#   def stack(self):
#   def stack(self):
#   def stack(self):
#   def __init__(self,
#                num_channels,
#                width_multiplier,
#                strides,
#                use_projection,
#                kernel_initializer=tf.keras.initializers.he_normal(),
#                kernel_regularizer=tf.keras.regularizers.l2(l=0.0001),
#                name='basic_residual_block'):
#   def compute_output_shape(self, input_shape):
#   def call(self, inputs, training):
#   def call(self, inputs, training):
#   def __init__(self,
#                num_channels,
#                width_multiplier,
#                strides,
#                use_projection,
#                kernel_initializer=tf.keras.initializers.he_normal(),
#                kernel_regularizer=tf.keras.regularizers.l2(l=0.0001),
#                name='bottleneck_residual_block'):
#   def compute_output_shape(self, input_shape):
#   def call(self, inputs, training):
#   def call(self, inputs, training):
#   def __init__(self,
#                kernel_initializer=tf.keras.initializers.he_normal(),
#                kernel_regularizer=tf.keras.regularizers.l2(l=0.0001),
#                large_input_kernel=True,
#                name='ResNet',
#                **kwargs):
#   def stack(self):
# def weight_standardization_replacements(model):
#   def call(self, inputs, **kwargs):
#   def __init__(self, width_multiplier, name='WideResNet', **kwargs):


import gin.tf
import tensorflow as tf
from meta_dataset.learners.experimental import base as learner_base
from meta_dataset.models.experimental import reparameterizable_backbones
      class_sums,
      tf.reshape(class_num_images, [-1] + [1] * (len(embeddings.shape) - 1)))

  return prototypes


@gin.configurable
def euclidean_distance(x, y):
  """Computes the Euclidean distance."""
  x = tf.expand_dims(x, 1)
  y = tf.expand_dims(y, 0)
  return tf.reduce_sum(input_tensor=tf.square(x - y), axis=2)


@gin.configurable
def exact_cosine_distance(x, y):
  """Computes the exact cosine distance."""
  x = tf.nn.l2_normalize(x, 1, epsilon=1e-3)
  y = tf.nn.l2_normalize(y, 1, epsilon=1e-3)
  return 1 - tf.matmul(x, y, transpose_b=True)


@gin.configurable
def partial_cosine_distance(x, y):
  """Computes the cosine distance without L2 normalization of x."""
  y = tf.nn.l2_normalize(y, 1, epsilon=1e-3)
  return 1 - tf.matmul(x, y, transpose_b=True)


@gin.configurable

# Path: meta_dataset/learners/experimental/base.py
# class NotBuiltError(RuntimeError):
# class ExperimentalLearner(learner_base.Learner):
# class ExperimentalEpisodicLearner(ExperimentalLearner,
#                                   learner_base.EpisodicLearner):
# class ExperimentalBatchLearner(ExperimentalLearner, learner_base.BatchLearner):
#   def __init__(self):
# def class_specific_data(onehot_labels, data, num_classes, axis=0):
#   def __init__(self, **kwargs):
#   def compute_regularizer(self, onehot_labels, predictions):
#   def build(self):
#   def variables(self):
#   def trainable_variables(self):
# 
# Path: meta_dataset/models/experimental/reparameterizable_backbones.py
# class ReparameterizableBackbone(tf.keras.Sequential,
#                                 reparameterizable_base.ReparameterizableModule):
# class ReparameterizableSpatialBackbone(ReparameterizableBackbone):
# class LinearModel(ReparameterizableBackbone):
# class FullyConnectedNet(ReparameterizableBackbone):
# class ConvNet(ReparameterizableSpatialBackbone):
# class RelationNetConvNet(ReparameterizableSpatialBackbone):
# class RelationModule(ReparameterizableBackbone):
# class BasicBlock(tf.keras.Model):
# class BasicBlockV1(BasicBlock):
# class BasicBlockV2(BasicBlock):
# class BottleneckBlock(tf.keras.Model):
# class BottleneckBlockV1(BottleneckBlock):
# class BottleneckBlockV2(BottleneckBlock):
# class ResNet(ReparameterizableSpatialBackbone):
# class ResNet18(ResNet):
# class ResNet18V1(ResNet18):
# class ResNet18V2(ResNet18):
# class ResNet34(ResNet):
# class ResNet34V1(ResNet34):
# class ResNet34V2(ResNet34):
# class ResNet50(ResNet):
# class ResNet50V1(ResNet50):
# class ResNet50V2(ResNet50):
# class ResNet101(ResNet):
# class ResNet101V1(ResNet101):
# class ResNet101V2(ResNet101):
# class ResNet152(ResNet):
# class ResNet152V1(ResNet152):
# class ResNet152V2(ResNet152):
# class GNResNet(ResNet):
# class GNBasicBlockV2(BasicBlockV2):
# class GNResNet18V2(GNResNet):
# class GNBottleneckBlockV2(BottleneckBlockV2):
# class GNResNet50V2(GNResNet):
# class WideResNet(ResNet):
# class WideResNet16(WideResNet):
# class WideResNet16V1(WideResNet16):
# class WideResNet16V2(WideResNet16):
# class WideResNet28(WideResNet):
# class WideResNet28V1(WideResNet28):
# class WideResNet28V2(WideResNet28):
#   def __init__(self,
#                output_dim=None,
#                kernel_initializer=tf.keras.initializers.he_normal(),
#                kernel_regularizer=tf.keras.regularizers.l2(),
#                parameter_adapter=None,
#                name=None):
#   def build(self, input_shape):
#   def stack(self):
#   def output_layers(self):
#   def __init__(self, keep_spatial_dims=False, **kwargs):
#   def output_layers(self):
#   def __init__(self, output_dim, **kwargs):
#   def stack(self):
#   def __init__(self, num_hiddens_per_layer, **kwargs):
#   def stack(self):
#   def __init__(self, num_filters_per_layer, **kwargs):
#   def stack(self):
#   def stack(self):
#   def stack(self):
#   def __init__(self,
#                num_channels,
#                width_multiplier,
#                strides,
#                use_projection,
#                kernel_initializer=tf.keras.initializers.he_normal(),
#                kernel_regularizer=tf.keras.regularizers.l2(l=0.0001),
#                name='basic_residual_block'):
#   def compute_output_shape(self, input_shape):
#   def call(self, inputs, training):
#   def call(self, inputs, training):
#   def __init__(self,
#                num_channels,
#                width_multiplier,
#                strides,
#                use_projection,
#                kernel_initializer=tf.keras.initializers.he_normal(),
#                kernel_regularizer=tf.keras.regularizers.l2(l=0.0001),
#                name='bottleneck_residual_block'):
#   def compute_output_shape(self, input_shape):
#   def call(self, inputs, training):
#   def call(self, inputs, training):
#   def __init__(self,
#                kernel_initializer=tf.keras.initializers.he_normal(),
#                kernel_regularizer=tf.keras.regularizers.l2(l=0.0001),
#                large_input_kernel=True,
#                name='ResNet',
#                **kwargs):
#   def stack(self):
# def weight_standardization_replacements(model):
#   def call(self, inputs, **kwargs):
#   def __init__(self, width_multiplier, name='WideResNet', **kwargs):


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

      class_sums,
      tf.reshape(class_num_images, [-1] + [1] * (len(embeddings.shape) - 1)))

  return prototypes


@gin.configurable
def euclidean_distance(x, y):
  """Computes the Euclidean distance."""
  x = tf.expand_dims(x, 1)
  y = tf.expand_dims(y, 0)
  return tf.reduce_sum(input_tensor=tf.square(x - y), axis=2)


@gin.configurable
def exact_cosine_distance(x, y):
  """Computes the exact cosine distance."""
  x = tf.nn.l2_normalize(x, 1, epsilon=1e-3)
  y = tf.nn.l2_normalize(y, 1, epsilon=1e-3)
  return 1 - tf.matmul(x, y, transpose_b=True)


@gin.configurable
def partial_cosine_distance(x, y):
  """Computes the cosine distance without L2 normalization of x."""
  y = tf.nn.l2_normalize(y, 1, epsilon=1e-3)
  return 1 - tf.matmul(x, y, transpose_b=True)


@gin.configurable

class MetricLearner(learner_base.ExperimentalEpisodicLearner):