# Path: meta_dataset/test_utils.py
# def get_argspec(fn):
# def get_inherited_args(cls, cls_method):
# def get_valid_kwargs(module_cls, valid_module_init_args,
#                      valid_module_call_args):
# 
# 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 numpy as np
import tensorflow.compat.v1 as tf
from absl.testing import parameterized
from meta_dataset import test_utils
from meta_dataset.models.experimental import reparameterizable_base

VALID_MODULE_INPUT_SHAPE = {
    ReparameterizableDense: (1, 32),
    ReparameterizableConv2D: (1, 32, 32, 3),
    ReparameterizableBatchNormalization: (1, 32),
}

# Transformations with constant Jacobian (e.g., linear transforms) should
# produce the same output gradient when variables are replaced.
# Note that transforms in this list may take an `activation` parameter
# that when set to a non-linearity such as `tf.nn.relu`, induces a non-constant
# Jacobian.
CONSTANT_JACOBIAN_TRANSFORMS = (
    ReparameterizableDense,
    ReparameterizableConv2D,
    ReparameterizableBatchNormalization,
)


def has_constant_jacobian(layer):
  return (any(isinstance(layer, cls) for cls in CONSTANT_JACOBIAN_TRANSFORMS)
          and (not hasattr(layer, 'activation') or layer.activation is None or
               'linear' in str(layer.activation)))


def _get_module_predicate_test_cases(module_cls, variable_predicate,
                                     module_predicate, valid_module_init_args,
                                     valid_module_call_args):
  """Return parameters of tests for `module_cls` with the given arguments."""
  test_cases = []

# Path: meta_dataset/test_utils.py
# def get_argspec(fn):
# def get_inherited_args(cls, cls_method):
# def get_valid_kwargs(module_cls, valid_module_init_args,
#                      valid_module_call_args):
# 
# 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_base_test.py
import itertools
import numpy as np
import tensorflow.compat.v1 as tf
from absl.testing import parameterized
from meta_dataset import test_utils
from meta_dataset.models.experimental import reparameterizable_base


VALID_MODULE_INPUT_SHAPE = {
    ReparameterizableDense: (1, 32),
    ReparameterizableConv2D: (1, 32, 32, 3),
    ReparameterizableBatchNormalization: (1, 32),
}

# Transformations with constant Jacobian (e.g., linear transforms) should
# produce the same output gradient when variables are replaced.
# Note that transforms in this list may take an `activation` parameter
# that when set to a non-linearity such as `tf.nn.relu`, induces a non-constant
# Jacobian.
CONSTANT_JACOBIAN_TRANSFORMS = (
    ReparameterizableDense,
    ReparameterizableConv2D,
    ReparameterizableBatchNormalization,
)


def has_constant_jacobian(layer):
  return (any(isinstance(layer, cls) for cls in CONSTANT_JACOBIAN_TRANSFORMS)
          and (not hasattr(layer, 'activation') or layer.activation is None or
               'linear' in str(layer.activation)))


def _get_module_predicate_test_cases(module_cls, variable_predicate,
                                     module_predicate, valid_module_init_args,
                                     valid_module_call_args):
  """Return parameters of tests for `module_cls` with the given arguments."""
  test_cases = []

  for valid_init_kwargs, valid_call_kwargs in test_utils.get_valid_kwargs(