1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
|
import tensorflow.python.platform
import numpy
import tensorflow as tf
class GenerateIdentityTensorTest(tf.test.TestCase):
def _testDiagOp(self, diag, dtype, expected_ans, use_gpu=False,
expected_err_re=None):
with self.test_session(use_gpu=use_gpu):
tf_ans = tf.diag(tf.convert_to_tensor(diag.astype(dtype)))
out = tf_ans.eval()
self.assertAllClose(out, expected_ans)
self.assertShapeEqual(expected_ans, tf_ans)
def testEmptyTensor(self):
x = numpy.array([])
expected_ans = numpy.empty([0, 0])
self._testDiagOp(x, numpy.int32, expected_ans)
def testRankOneIntTensor(self):
x = numpy.array([1, 2, 3])
expected_ans = numpy.array(
[[1, 0, 0],
[0, 2, 0],
[0, 0, 3]])
self._testDiagOp(x, numpy.int32, expected_ans)
self._testDiagOp(x, numpy.int64, expected_ans)
def testRankOneFloatTensor(self):
x = numpy.array([1.1, 2.2, 3.3])
expected_ans = numpy.array(
[[1.1, 0, 0],
[0, 2.2, 0],
[0, 0, 3.3]])
self._testDiagOp(x, numpy.float32, expected_ans)
self._testDiagOp(x, numpy.float64, expected_ans)
def testRankTwoIntTensor(self):
x = numpy.array([[1, 2, 3], [4, 5, 6]])
expected_ans = numpy.array(
[[[[1, 0, 0], [0, 0, 0]],
[[0, 2, 0], [0, 0, 0]],
[[0, 0, 3], [0, 0, 0]]],
[[[0, 0, 0], [4, 0, 0]],
[[0, 0, 0], [0, 5, 0]],
[[0, 0, 0], [0, 0, 6]]]])
self._testDiagOp(x, numpy.int32, expected_ans)
self._testDiagOp(x, numpy.int64, expected_ans)
def testRankTwoFloatTensor(self):
x = numpy.array([[1.1, 2.2, 3.3], [4.4, 5.5, 6.6]])
expected_ans = numpy.array(
[[[[1.1, 0, 0], [0, 0, 0]],
[[0, 2.2, 0], [0, 0, 0]],
[[0, 0, 3.3], [0, 0, 0]]],
[[[0, 0, 0], [4.4, 0, 0]],
[[0, 0, 0], [0, 5.5, 0]],
[[0, 0, 0], [0, 0, 6.6]]]])
self._testDiagOp(x, numpy.float32, expected_ans)
self._testDiagOp(x, numpy.float64, expected_ans)
def testRankThreeFloatTensor(self):
x = numpy.array([[[1.1, 2.2], [3.3, 4.4]],
[[5.5, 6.6], [7.7, 8.8]]])
expected_ans = numpy.array(
[[[[[[1.1, 0], [0, 0]], [[0, 0], [0, 0]]],
[[[0, 2.2], [0, 0]], [[0, 0], [0, 0]]]],
[[[[0, 0], [3.3, 0]], [[0, 0], [0, 0]]],
[[[0, 0], [0, 4.4]], [[0, 0], [0, 0]]]]],
[[[[[0, 0], [0, 0]], [[5.5, 0], [0, 0]]],
[[[0, 0], [0, 0]], [[0, 6.6], [0, 0]]]],
[[[[0, 0], [0, 0]], [[0, 0], [7.7, 0]]],
[[[0, 0], [0, 0]], [[0, 0], [0, 8.8]]]]]])
self._testDiagOp(x, numpy.float32, expected_ans)
self._testDiagOp(x, numpy.float64, expected_ans)
if __name__ == "__main__":
tf.test.main()
|
