# 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
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Tests for relational_layers.py."""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function



def _create_positional_encoding_matrices():
  """Shared input/output pair for the positional encoding tests."""
  input_array = np.arange(24, dtype=np.float64).reshape((1, 4, 3, 2))
  output_array = np.eye(4)
  output_array = np.repeat(np.expand_dims(output_array, -1), 2, axis=-1)
  output_array = np.expand_dims(output_array, 0)
  return input_array, output_array


class RelationalLayersTest(tf.test.TestCase):

  def test_repeat_for_tensor(self):
    a = np.arange(24).reshape((1, 4, 3, 2))
    shouldbe = np.concatenate([a] * 3, axis=-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

#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Tests for relational_layers.py."""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function



def _create_positional_encoding_matrices():
  """Shared input/output pair for the positional encoding tests."""
  input_array = np.arange(24, dtype=np.float64).reshape((1, 4, 3, 2))
  output_array = np.eye(4)
  output_array = np.repeat(np.expand_dims(output_array, -1), 2, axis=-1)
  output_array = np.expand_dims(output_array, 0)
  return input_array, output_array


class RelationalLayersTest(tf.test.TestCase):

  def test_repeat_for_tensor(self):
    a = np.arange(24).reshape((1, 4, 3, 2))
    shouldbe = np.concatenate([a] * 3, axis=-2)

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