# 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
import numpy as np
import tensorflow.compat.v1 as tf
    result = self.evaluate(relational_layers.repeat(tf.constant(a), 3, axis=-2))
    self.assertAllClose(shouldbe, result)

  def test_pairwise_edge_embeddings_for_tensor(self):
    a = np.array([[[1], [2]]])
    shouldbe = np.array([[[[1, 1], [1, 2]], [[2, 1], [2, 2]]]])
    layer = relational_layers.PairwiseEdgeEmbeddings()
    result = self.evaluate(layer(tf.constant(a)))
    self.assertAllClose(shouldbe, result)

  def test_relational_layer_for_tensor(self):
    a = np.array([[[1], [2]]])
    shouldbe = np.array([[[2, 3], [4, 3]]])
    layer = relational_layers.RelationalLayer(
        tf.keras.layers.Lambda(lambda x: x),
        tf.keras.layers.Lambda(lambda x: tf.reduce_sum(x, axis=-2)))
    result = self.evaluate(layer(tf.constant(a)))
    self.assertAllClose(shouldbe, result)

  def test_positional_encoding_like_for_static_shape_tensor(self):
    value, shouldbe = _create_positional_encoding_matrices()
    a = tf.constant(value)
    output_tensor = relational_layers.positional_encoding_like(a, -3, -2)
    result = self.evaluate(output_tensor)
    self.assertEqual((1, 4, 4, 2), result.shape)
    self.assertAllClose(shouldbe, result)

  def test_positional_encoding_like_for_dynamic_shape_tensor(self):
    value, shouldbe = _create_positional_encoding_matrices()
    a = tf.placeholder(tf.float32, shape=(None, 4, 3, 2))

# 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/relational_layers_test.py
from disentanglement_lib.evaluation.abstract_reasoning import relational_layers
import numpy as np
import tensorflow.compat.v1 as tf

    result = self.evaluate(relational_layers.repeat(tf.constant(a), 3, axis=-2))
    self.assertAllClose(shouldbe, result)

  def test_pairwise_edge_embeddings_for_tensor(self):
    a = np.array([[[1], [2]]])
    shouldbe = np.array([[[[1, 1], [1, 2]], [[2, 1], [2, 2]]]])
    layer = relational_layers.PairwiseEdgeEmbeddings()
    result = self.evaluate(layer(tf.constant(a)))
    self.assertAllClose(shouldbe, result)

  def test_relational_layer_for_tensor(self):
    a = np.array([[[1], [2]]])
    shouldbe = np.array([[[2, 3], [4, 3]]])
    layer = relational_layers.RelationalLayer(
        tf.keras.layers.Lambda(lambda x: x),
        tf.keras.layers.Lambda(lambda x: tf.reduce_sum(x, axis=-2)))
    result = self.evaluate(layer(tf.constant(a)))
    self.assertAllClose(shouldbe, result)

  def test_positional_encoding_like_for_static_shape_tensor(self):
    value, shouldbe = _create_positional_encoding_matrices()
    a = tf.constant(value)
    output_tensor = relational_layers.positional_encoding_like(a, -3, -2)
    result = self.evaluate(output_tensor)
    self.assertEqual((1, 4, 4, 2), result.shape)
    self.assertAllClose(shouldbe, result)

  def test_positional_encoding_like_for_dynamic_shape_tensor(self):
    value, shouldbe = _create_positional_encoding_matrices()
    a = tf.placeholder(tf.float32, shape=(None, 4, 3, 2))

    output_tensor = relational_layers.positional_encoding_like(a, -3, -2)