def test_multi_input(self):
op = LambdaOp(fn=np.reshape)
data = op.forward(data=[np.array([1, 2, 3, 4]), (2, 2)], state={})
self.assertTrue(is_equal(data, np.array([[1, 2], [3, 4]])))
def test_single_input(self):
op = LambdaOp(fn=np.sum)
data = op.forward(data=[[1, 2, 3]], state={})
self.assertEqual(data, 6)