# Path: disentanglement_lib/evaluation/abstract_reasoning/relational_layers.py
# class RelationalLayer(tf.keras.layers.Layer):
# class PairwiseEdgeEmbeddings(tf.keras.layers.Layer):
# class AddPositionalEncoding(tf.keras.layers.Layer):
# class StackAnswers(tf.keras.layers.Layer):
# class MultiDimBatchApply(tf.keras.layers.Layer):
#   def __init__(self, edge_layer, reduce_layer):
#   def call(self, inputs, **kwargs):
#   def call(self, inputs, **kwargs):
# def repeat(tensor, num, axis):
# def positional_encoding_like(tensor, positional_encoding_axis=-2,
#                              value_axis=-1):
#   def __init__(self, positional_encoding_axis=-2, embedding_axis=-1):
#   def call(self, inputs, **kwargs):
#   def __init__(self, answer_axis=-2, stack_axis=-3):
#   def call(self, inputs, **kwargs):
#   def __init__(self, layer, num_dims_to_keep=3):
#   def call(self, inputs, **kwargs):


from disentanglement_lib.evaluation.abstract_reasoning import relational_layers
from tensorflow.contrib import tpu as contrib_tpu
import gin
import tensorflow.compat.v1 as tf
import tensorflow_hub as hub
      **kwargs: Other keyword arguments passed to tf.keras.Model.
    """
    super(BaselineCNNEmbedder, self).__init__(name=name, **kwargs)
    embedding_layers = [
        tf.keras.layers.Conv2D(
            32, (4, 4),
            2,
            activation=get_activation(),
            padding="same",
            kernel_initializer=get_kernel_initializer()),
        tf.keras.layers.Conv2D(
            32, (4, 4),
            2,
            activation=get_activation(),
            padding="same",
            kernel_initializer=get_kernel_initializer()),
        tf.keras.layers.Conv2D(
            64, (4, 4),
            2,
            activation=get_activation(),
            padding="same",
            kernel_initializer=get_kernel_initializer()),
        tf.keras.layers.Conv2D(
            64, (4, 4),
            2,
            activation=get_activation(),
            padding="same",
            kernel_initializer=get_kernel_initializer()),
        tf.keras.layers.Flatten(),
    ]

# Path: disentanglement_lib/evaluation/abstract_reasoning/relational_layers.py
# class RelationalLayer(tf.keras.layers.Layer):
# class PairwiseEdgeEmbeddings(tf.keras.layers.Layer):
# class AddPositionalEncoding(tf.keras.layers.Layer):
# class StackAnswers(tf.keras.layers.Layer):
# class MultiDimBatchApply(tf.keras.layers.Layer):
#   def __init__(self, edge_layer, reduce_layer):
#   def call(self, inputs, **kwargs):
#   def call(self, inputs, **kwargs):
# def repeat(tensor, num, axis):
# def positional_encoding_like(tensor, positional_encoding_axis=-2,
#                              value_axis=-1):
#   def __init__(self, positional_encoding_axis=-2, embedding_axis=-1):
#   def call(self, inputs, **kwargs):
#   def __init__(self, answer_axis=-2, stack_axis=-3):
#   def call(self, inputs, **kwargs):
#   def __init__(self, layer, num_dims_to_keep=3):
#   def call(self, inputs, **kwargs):


# Path: disentanglement_lib/evaluation/abstract_reasoning/models.py
from disentanglement_lib.evaluation.abstract_reasoning import relational_layers
from tensorflow.contrib import tpu as contrib_tpu
import gin
import tensorflow.compat.v1 as tf
import tensorflow_hub as hub

      **kwargs: Other keyword arguments passed to tf.keras.Model.
    """
    super(BaselineCNNEmbedder, self).__init__(name=name, **kwargs)
    embedding_layers = [
        tf.keras.layers.Conv2D(
            32, (4, 4),
            2,
            activation=get_activation(),
            padding="same",
            kernel_initializer=get_kernel_initializer()),
        tf.keras.layers.Conv2D(
            32, (4, 4),
            2,
            activation=get_activation(),
            padding="same",
            kernel_initializer=get_kernel_initializer()),
        tf.keras.layers.Conv2D(
            64, (4, 4),
            2,
            activation=get_activation(),
            padding="same",
            kernel_initializer=get_kernel_initializer()),
        tf.keras.layers.Conv2D(
            64, (4, 4),
            2,
            activation=get_activation(),
            padding="same",
            kernel_initializer=get_kernel_initializer()),
        tf.keras.layers.Flatten(),
    ]

    self.embedding_layer = relational_layers.MultiDimBatchApply(