Skip to content

Commit

Permalink
[Zero-Dim] Support input 0D Tensor for masked_select (#49803)
Browse files Browse the repository at this point in the history
* [Zero-Dim] Support input 0D Tensor for masked_select
  • Loading branch information
ZzSean authored Jan 17, 2023
1 parent 2242136 commit ce04589
Show file tree
Hide file tree
Showing 2 changed files with 44 additions and 0 deletions.
30 changes: 30 additions & 0 deletions python/paddle/fluid/tests/unittests/test_zero_dim_tensor.py
Original file line number Diff line number Diff line change
Expand Up @@ -1298,6 +1298,20 @@ def test_atan2(self):
self.assertEqual(x1.grad.numpy(), 0.5)
self.assertEqual(x2.grad.numpy(), 0)

def test_maseked_select(self):
x = paddle.rand([])
x.stop_gradient = False
mask = paddle.full([], True, dtype='bool')
y = paddle.masked_select(x, mask)

y.retain_grads()
y.backward()
self.assertEqual(y.shape, [1])
self.assertEqual(y.numpy(), x.numpy())
self.assertEqual(y.grad.shape, [1])
self.assertEqual(x.grad.shape, [])
self.assertEqual(x.grad.numpy(), 1)


class TestSundryAPIStatic(unittest.TestCase):
def setUp(self):
Expand Down Expand Up @@ -1968,6 +1982,22 @@ def test_atan2(self):

self.assertEqual(res[0].shape, ())

@prog_scope()
def test_maseked_select(self):
x = paddle.rand([])
x.stop_gradient = False
mask = paddle.full([], True, dtype='bool')
y = paddle.masked_select(x, mask)
paddle.static.append_backward(y.sum())

prog = paddle.static.default_main_program()
res = self.exe.run(prog, fetch_list=[x, y, y.grad_name, x.grad_name])
self.assertEqual(res[1].shape, (1,))
self.assertEqual(res[1], res[0])
self.assertEqual(res[2].shape, (1,))
self.assertEqual(res[3].shape, ())
self.assertEqual(res[3], 1)


# Use to test API whose zero-dim input tensors don't have grad and not need to test backward in OpTest.
class TestNoBackwardAPI(unittest.TestCase):
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -787,6 +787,20 @@ def test_equalall(self):
self.assertEqual(out.shape, [])
self.assertFalse(out)

def test_maseked_select(self):
x = paddle.rand([])
x.stop_gradient = False
mask = paddle.full([], True, dtype='bool')
y = paddle.masked_select(x, mask)

y.retain_grads()
y.backward()
self.assertEqual(y.shape, [1])
self.assertEqual(y.numpy(), x.numpy())
self.assertEqual(y.grad.shape, [1])
self.assertEqual(x.grad.shape, [])
self.assertEqual(x.grad.numpy(), 1)


# Use to test API whose zero-dim input tensors don't have grad and not need to test backward in OpTest.

Expand Down

0 comments on commit ce04589

Please sign in to comment.