def setUp(self):
self.params = Parameters()
self.model_pb = Model()
self.tensors_pb = self.model_pb.param
self.embeddings_pb = self.model_pb.embedding_table_info
arr1 = np.random.uniform(size=(3, 4))
tensor1_pb = Tensor(arr1, name="x").to_tensor_pb()
arr2 = np.random.uniform(size=(4, 5))
tensor2_pb = Tensor(arr2, name="y").to_tensor_pb()
self.tensors_pb.extend([tensor1_pb, tensor2_pb])
self.embedding_table_name = "embedding_1"
self.embedding_dim = 10
embedding_pb = EmbeddingTableInfo()
embedding_pb.name = self.embedding_table_name
embedding_pb.dim = self.embedding_dim
embedding_pb.initializer = "uniform"
embedding_vectors = np.random.uniform(size=(2, 10))
embedding_indices = np.array([0, 8])
embedding_tensor = Tensor(
embedding_vectors,
indices=embedding_indices,
name=self.embedding_table_name,
)
embedding_tensor_pb = embedding_tensor.to_tensor_pb()
self.tensors_pb.append(embedding_tensor_pb)
self.embeddings_pb.append(embedding_pb)
def test_tensor_data_structure(self):
# Test tensor values, without indices
arr = np.ndarray(shape=[3, 1, 2, 4], dtype=np.int32)
tensor = Tensor(arr)
self.assertTrue(np.array_equal(arr, tensor.values))
self.assertTrue(np.array_equal(arr, tensor.to_tf_tensor()))
self.assertFalse(tensor.is_indexed_slices())
# Test tensor values, with indices
indices = np.array([2, 0, 1])
tensor = Tensor(arr, indices)
self.assertTrue(np.array_equal(arr, tensor.values))
self.assertTrue(np.array_equal(indices, tensor.indices))
self.assertTrue(np.array_equal(arr, tensor.to_tf_tensor().values))
self.assertTrue(np.array_equal(indices, tensor.to_tf_tensor().indices))
self.assertTrue(tensor.is_indexed_slices())
# Test round trip
# tensor to tensor PB
tensor = Tensor(arr, indices, name="test")
pb = tensor.to_tensor_pb()
self.assertEqual(pb.name, "test")
self.assertEqual(pb.dim, [3, 1, 2, 4])
self.assertEqual(pb.dtype, tensor_dtype_pb2.DT_INT32)
np.testing.assert_array_equal(pb.indices, indices)
# tensor PB to tensor
tensor_new = Tensor.from_tensor_pb(pb)
self.assertEqual(tensor.name, "test")
np.testing.assert_array_equal(tensor_new.values, arr)
np.testing.assert_array_equal(tensor_new.indices, indices)
# Test Tensor().to_ndarray()
values = np.array([[1.0, 2.0], [3.0, 4.0]])
indices = np.array([0, 2])
name = "test"
tensor = Tensor(values, indices, name)
self.assertRaises(NotImplementedError, tensor.to_ndarray)
tensor = Tensor(values, name=name)
self.assertTrue(np.allclose(values, tensor.to_ndarray()))