# 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_base.py
# def is_variable(obj):
# def is_trainable_variable(obj):
# def is_batch_norm_module(obj):
# def corner_case_setattr(target, attr, value):
# def corner_case_getattr(target, attr):
# def chained_getattr(obj, path):
# def _patch_getattribute(cls, new_getattribute):
# def _custom_getter(getter, classes=(tf.Module,)):
#     def new_getattribute(obj, name, orig_getattribute=orig_getattribute):
#   def reparameterize(self, replacements):
#     def getter(v_source):
#   def reparameterizables(self,
#                          variable_predicate=is_variable,
#                          module_predicate=None,
#                          with_path=False):
#       def path_prefixes(path):
#       def satisfies_module_predicate(variable_with_path):
# class ReparameterizableModule(tf.Module):


import itertools
import math
import gin.tf
import tensorflow as tf
import tensorflow_probability as tfp
from meta_dataset.learners.experimental import base as learner_base
from meta_dataset.models.experimental import reparameterizable_base
  def __init__(self, estimate_scale, damping, **kwargs):
    """A class mixture.

    Args:
      estimate_scale: Whether to estimate the scale in addition to the location.
      damping: The scale of the covariance damping coefficient.
      **kwargs: Keyword arguments common to `ReparameterizableClassMixture`s.
    """
    super(MultivariateNormalDiag, self).__init__(**kwargs)
    self.damping = damping
    self.estimate_scale = estimate_scale

  def build_meta_parameters(self):
    """Assign to attributes the task parameters."""
    return

  def build_task_parameters(self):
    """Assign to attributes the meta parameters."""
    self.locs = [
        tf.Variable(tf.zeros((self.num_dims)), name='loc_{}'.format(i))
        for i in range(self.num_components)
    ]
    self.log_scales = [
        tf.Variable(tf.zeros((self.num_dims)), name='log_scale_{}'.format(i))
        for i in range(self.num_components)
    ]

  def episodic_init_ops(self, onehot_labels, embeddings, task_parameters):
    """Perform (data-dependent) initialization operations."""
    del task_parameters

# 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_base.py
# def is_variable(obj):
# def is_trainable_variable(obj):
# def is_batch_norm_module(obj):
# def corner_case_setattr(target, attr, value):
# def corner_case_getattr(target, attr):
# def chained_getattr(obj, path):
# def _patch_getattribute(cls, new_getattribute):
# def _custom_getter(getter, classes=(tf.Module,)):
#     def new_getattribute(obj, name, orig_getattribute=orig_getattribute):
#   def reparameterize(self, replacements):
#     def getter(v_source):
#   def reparameterizables(self,
#                          variable_predicate=is_variable,
#                          module_predicate=None,
#                          with_path=False):
#       def path_prefixes(path):
#       def satisfies_module_predicate(variable_with_path):
# class ReparameterizableModule(tf.Module):


# Path: meta_dataset/models/experimental/reparameterizable_distributions.py
import itertools
import math
import gin.tf
import tensorflow as tf
import tensorflow_probability as tfp
from meta_dataset.learners.experimental import base as learner_base
from meta_dataset.models.experimental import reparameterizable_base

  def __init__(self, estimate_scale, damping, **kwargs):
    """A class mixture.

    Args:
      estimate_scale: Whether to estimate the scale in addition to the location.
      damping: The scale of the covariance damping coefficient.
      **kwargs: Keyword arguments common to `ReparameterizableClassMixture`s.
    """
    super(MultivariateNormalDiag, self).__init__(**kwargs)
    self.damping = damping
    self.estimate_scale = estimate_scale

  def build_meta_parameters(self):
    """Assign to attributes the task parameters."""
    return

  def build_task_parameters(self):
    """Assign to attributes the meta parameters."""
    self.locs = [
        tf.Variable(tf.zeros((self.num_dims)), name='loc_{}'.format(i))
        for i in range(self.num_components)
    ]
    self.log_scales = [
        tf.Variable(tf.zeros((self.num_dims)), name='log_scale_{}'.format(i))
        for i in range(self.num_components)
    ]

  def episodic_init_ops(self, onehot_labels, embeddings, task_parameters):
    """Perform (data-dependent) initialization operations."""
    del task_parameters

    class_embeddings = learner_base.class_specific_data(onehot_labels,