google-research

1
# coding=utf-8
2
# Copyright 2024 The Google Research Authors.
3
#
4
# Licensed under the Apache License, Version 2.0 (the "License");
5
# you may not use this file except in compliance with the License.
6
# You may obtain a copy of the License at
7
#
8
#     http://www.apache.org/licenses/LICENSE-2.0
9
#
10
# Unless required by applicable law or agreed to in writing, software
11
# distributed under the License is distributed on an "AS IS" BASIS,
12
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13
# See the License for the specific language governing permissions and
14
# limitations under the License.
15

16
"""Tests for pairs."""
17
import tensorflow as tf
18

19
from dedal import pairs as pairs_lib
20

21

22
class PairsTest(tf.test.TestCase):
23

24
  def setUp(self):
25
    super().setUp()
26
    tf.random.set_seed(0)
27

28
  def test_build(self):
29
    batch_size = 32
30
    dim = 5
31
    inputs = tf.random.uniform((batch_size * 2, dim))
32
    indices = tf.reshape(tf.range(batch_size *2), (-1, 2))
33

34
    result, = pairs_lib.build(indices, inputs)
35
    self.assertAllEqual(result[4], tf.stack([inputs[8], inputs[9]]))
36

37
    # Shift by one the second column.
38
    indices = tf.stack(
39
        [indices[:, 0], tf.roll(indices[:, 1], shift=1, axis=0)], axis=1)
40
    result, = pairs_lib.build(indices, inputs)
41
    self.assertAllEqual(result[4], tf.stack([inputs[8], inputs[7]]))
42

43

44
if __name__ == '__main__':
45
  tf.test.main()
46