|
|
|
@ -34,8 +34,8 @@ class TestSpliteSelectedRows(unittest.TestCase):
|
|
|
|
|
|
|
|
|
|
def check_with_place(self, place):
|
|
|
|
|
scope = core.Scope()
|
|
|
|
|
rows = [0, 5, 7, 4]
|
|
|
|
|
height = 10
|
|
|
|
|
rows = [0, 5, 7, 4, 20]
|
|
|
|
|
height = 20
|
|
|
|
|
row_numel = 2
|
|
|
|
|
|
|
|
|
|
# initialize input variable X
|
|
|
|
@ -45,38 +45,41 @@ class TestSpliteSelectedRows(unittest.TestCase):
|
|
|
|
|
np_array = np.ones((len(rows), row_numel)).astype("float32")
|
|
|
|
|
np_array[0, 0] = 2.0
|
|
|
|
|
np_array[2, 1] = 4.0
|
|
|
|
|
np_array[4, 1] = 8.0
|
|
|
|
|
x_tensor = x.get_tensor()
|
|
|
|
|
x_tensor.set(np_array, place)
|
|
|
|
|
|
|
|
|
|
rows_sections = [2, 2]
|
|
|
|
|
height_sections = []
|
|
|
|
|
height_sections = [5, 5, 5, 5, 3]
|
|
|
|
|
|
|
|
|
|
# initialize output variables [out0, out1]
|
|
|
|
|
out0 = scope.var('out0').get_selected_rows()
|
|
|
|
|
out1 = scope.var('out1').get_selected_rows()
|
|
|
|
|
outs_name = ["out%d" % i for i in xrange(len(height_sections))]
|
|
|
|
|
outs = [
|
|
|
|
|
scope.var(var_name).get_selected_rows() for var_name in outs_name
|
|
|
|
|
]
|
|
|
|
|
|
|
|
|
|
# expected output selected rows
|
|
|
|
|
expected_out0_rows = [0, 5]
|
|
|
|
|
expected_out1_rows = [7, 4]
|
|
|
|
|
expected_height = height
|
|
|
|
|
expected_out0_rows = [0, 4]
|
|
|
|
|
expected_out1_rows = [5, 7]
|
|
|
|
|
expected_out4_rows = [20]
|
|
|
|
|
|
|
|
|
|
op = Operator(
|
|
|
|
|
"split_selected_rows",
|
|
|
|
|
X="X",
|
|
|
|
|
Out=["out0", "out1"],
|
|
|
|
|
rows_sections=rows_sections,
|
|
|
|
|
Out=outs_name,
|
|
|
|
|
height_sections=height_sections)
|
|
|
|
|
|
|
|
|
|
op.run(scope, place)
|
|
|
|
|
|
|
|
|
|
self.assertEqual(out0.rows(), expected_out0_rows)
|
|
|
|
|
self.assertEqual(out1.rows(), expected_out1_rows)
|
|
|
|
|
self.assertEqual(outs[0].rows(), expected_out0_rows)
|
|
|
|
|
self.assertEqual(outs[1].rows(), expected_out1_rows)
|
|
|
|
|
self.assertEqual(outs[4].rows(), expected_out4_rows)
|
|
|
|
|
|
|
|
|
|
self.assertEqual(out0.height(), expected_height)
|
|
|
|
|
self.assertEqual(out1.height(), expected_height)
|
|
|
|
|
self.assertEqual(outs[0].height(), height_sections[0])
|
|
|
|
|
self.assertEqual(outs[4].height(), height_sections[4])
|
|
|
|
|
|
|
|
|
|
self.assertAlmostEqual(2.0, np.array(out0.get_tensor())[0, 0])
|
|
|
|
|
self.assertAlmostEqual(4.0, np.array(out1.get_tensor())[0, 1])
|
|
|
|
|
self.assertAlmostEqual(2.0, np.array(outs[0].get_tensor())[0, 0])
|
|
|
|
|
self.assertAlmostEqual(4.0, np.array(outs[1].get_tensor())[1, 1])
|
|
|
|
|
self.assertAlmostEqual(8.0, np.array(outs[4].get_tensor())[0, 1])
|
|
|
|
|
|
|
|
|
|
def check_grad_with_place(self, place):
|
|
|
|
|
scope = core.Scope()
|
|
|
|
@ -84,8 +87,7 @@ class TestSpliteSelectedRows(unittest.TestCase):
|
|
|
|
|
row_numel = 2
|
|
|
|
|
|
|
|
|
|
# attr
|
|
|
|
|
rows_sections = [2, 2]
|
|
|
|
|
height_sections = []
|
|
|
|
|
height_sections = [5, 5]
|
|
|
|
|
|
|
|
|
|
# initialize input variable X
|
|
|
|
|
out0_grad = scope.var("out0@GRAD").get_selected_rows()
|
|
|
|
@ -112,7 +114,6 @@ class TestSpliteSelectedRows(unittest.TestCase):
|
|
|
|
|
"sum",
|
|
|
|
|
X=["out0@GRAD", "out1@GRAD"],
|
|
|
|
|
Out="X@GRAD",
|
|
|
|
|
rows_sections=rows_sections,
|
|
|
|
|
height_sections=height_sections)
|
|
|
|
|
|
|
|
|
|
grad_op.run(scope, place)
|
|
|
|
|