diff --git a/tensorflow_quantum/core/ops/tfq_utility_ops_test.py b/tensorflow_quantum/core/ops/tfq_utility_ops_test.py index 613704fe2..cda7cbe8b 100644 --- a/tensorflow_quantum/core/ops/tfq_utility_ops_test.py +++ b/tensorflow_quantum/core/ops/tfq_utility_ops_test.py @@ -139,6 +139,12 @@ def test_append_circuit(self, max_n_bits, symbols, n_circuits): 'padded_array': [[[0, 0, 0, 0], [1, 1, 1, 1]]] }, { 'padded_array': [[[1, 1, -2, -2], [0, 1, -2, -2], [0, 0, -2, -2]]] + }, { + 'padded_array': [[[-2, -2], [-2, -2]]] + }, { + 'padded_array': [[[1, 1], [1, 1]]] + }, { + 'padded_array': np.empty((0, 2, 2), dtype=np.float32) }]) def test_padded_to_ragged(self, padded_array): """Test for padded_to_ragged utility.""" @@ -148,6 +154,27 @@ def test_padded_to_ragged(self, padded_array): np.array(padded_array, dtype=float)) self.assertAllEqual(expected, actual) + @parameterized.parameters([{ + 'padded_array': [[[1, 0, -2], [0, 1, -2], [-2, -2, -2]], + [[1, -2, -2], [-2, -2, -2], [-2, -2, -2]]] + }, { + 'padded_array': [[[1, 0, 0], [0, 1, 0], [0, 0, 1]]] + }, { + 'padded_array': [[[-2, -2, -2], [-2, -2, -2], [-2, -2, -2]]] + }, { + 'padded_array': np.empty((0, 3, 3), dtype=np.float32) + }]) + def test_padded_to_ragged2d(self, padded_array): + """Test for padded_to_ragged2d utility.""" + tensor_arr = tf.constant(padded_array, dtype=tf.float32) + row_mask = tf.abs(tensor_arr[:, :, 0]) < 1.1 + masked_rows = tf.ragged.boolean_mask(tensor_arr, row_mask) + element_mask = tf.abs(masked_rows) < 1.1 + expected = tf.ragged.boolean_mask(masked_rows, element_mask) + actual = tfq_utility_ops.padded_to_ragged2d( + np.array(padded_array, dtype=float)) + self.assertAllEqual(expected, actual) + class ResolveParametersOpTest(tf.test.TestCase, parameterized.TestCase): """Test the in-graph parameter resolving op."""