diff --git a/python/paddle/fluid/tests/unittests/test_tril_triu_op.py b/python/paddle/fluid/tests/unittests/test_tril_triu_op.py index 84dc7bd8a9659898056cccec42e030f92417880d..0f14c9d1c3ba99b5c9b1500e4b7ddabb690e9290 100644 --- a/python/paddle/fluid/tests/unittests/test_tril_triu_op.py +++ b/python/paddle/fluid/tests/unittests/test_tril_triu_op.py @@ -134,6 +134,14 @@ class TestTrilTriuOpAPI(unittest.TestCase): self.assertTrue(np.allclose(tril_out, np.tril(data))) self.assertTrue(np.allclose(triu_out, np.triu(data))) + def test_api_with_dygraph(self): + with fluid.dygraph.guard(): + data = np.random.random([1, 9, 9, 4]).astype('float32') + x = fluid.dygraph.to_variable(data) + tril_out, triu_out = tensor.tril(x).numpy(), tensor.triu(x).numpy() + self.assertTrue(np.allclose(tril_out, np.tril(data))) + self.assertTrue(np.allclose(triu_out, np.triu(data))) + if __name__ == '__main__': unittest.main() diff --git a/python/paddle/tensor/creation.py b/python/paddle/tensor/creation.py index 09cb6bec788057e914b02bacd6095b470c3ad64e..28bf09175490bb36fefb318087b3179da32d3065 100644 --- a/python/paddle/tensor/creation.py +++ b/python/paddle/tensor/creation.py @@ -696,8 +696,6 @@ def tril(input, diagonal=0, name=None): # [ 5, 6, 0, 0], # [ 9, 10, 11, 0]]) - .. code-block:: python - # example 2, positive diagonal value tril = tensor.tril(x, diagonal=2) tril_out, = exe.run(fluid.default_main_program(), feed={"x": data}, @@ -706,8 +704,6 @@ def tril(input, diagonal=0, name=None): # [ 5, 6, 7, 8], # [ 9, 10, 11, 12]]) - .. code-block:: python - # example 3, negative diagonal value tril = tensor.tril(x, diagonal=-1) tril_out, = exe.run(fluid.default_main_program(), feed={"x": data}, @@ -716,7 +712,10 @@ def tril(input, diagonal=0, name=None): # [ 5, 0, 0, 0], # [ 9, 10, 0, 0]]) - """ + """ + if in_dygraph_mode(): + op = getattr(core.ops, 'tril_triu') + return op(input, 'diagonal', diagonal, "lower", True) return _tril_triu_op(LayerHelper('tril', **locals())) @@ -771,8 +770,6 @@ def triu(input, diagonal=0, name=None): # [ 0, 6, 7, 8], # [ 0, 0, 11, 12]]) - .. code-block:: python - # example 2, positive diagonal value triu = tensor.triu(x, diagonal=2) triu_out, = exe.run(fluid.default_main_program(), feed={"x": data}, @@ -781,8 +778,6 @@ def triu(input, diagonal=0, name=None): # [0, 0, 0, 8], # [0, 0, 0, 0]]) - .. code-block:: python - # example 3, negative diagonal value triu = tensor.triu(x, diagonal=-1) triu_out, = exe.run(fluid.default_main_program(), feed={"x": data}, @@ -792,6 +787,9 @@ def triu(input, diagonal=0, name=None): # [ 0, 10, 11, 12]]) """ + if in_dygraph_mode(): + op = getattr(core.ops, 'tril_triu') + return op(input, 'diagonal', diagonal, "lower", False) return _tril_triu_op(LayerHelper('triu', **locals()))