# Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. import unittest import paddle from paddle.fluid.layer_helper import LayerHelper from paddle.fluid.layers.utils import flatten from paddle.incubate.autograd.primrules import _orig2prim, _prim2orig, _jvp, _transpose paddle.enable_static() ############################ Test prim2orig rules ############################ class TestAddPPrim2Orig(unittest.TestCase): def setUp(self): self.main_program = paddle.static.Program() self.startup_program = paddle.static.Program() self.layer_help = LayerHelper('TestPrim2Orig') with paddle.static.program_guard(self.main_program, self.startup_program): self.init_data() def init_data(self): self.op_type = 'add_p' X = paddle.static.data(name='X', shape=[2, 2], dtype='float') Y = paddle.static.data(name='Y', shape=[2, 2], dtype='float') self.input = {'X': X, 'Y': Y} self.output = { 'Z': self.layer_help.create_variable_for_type_inference(dtype=X.dtype) } self.attrs = {} self.prim2orig_args = (X, Y) self.all_ops = ['add_p', 'elementwise_add'] # { prim_op_output_var: orign_op_out_index } self.out_map = {self.output['Z']: 0} def test_op(self): with paddle.static.program_guard(self.main_program, self.startup_program): op = self.layer_help.append_op(type=self.op_type, inputs=self.input, outputs=self.output, attrs=self.attrs) orig_out = _prim2orig(op, *self.prim2orig_args) all_ops = [op.type for op in self.main_program.block(0).ops] self.assertEqual(sorted(all_ops), sorted(self.all_ops)) orig_out = flatten(orig_out) for k, v in self.out_map.items(): self.assertEqual(k.shape, orig_out[v].shape) class TestSubPPrim2Orig(TestAddPPrim2Orig): def init_data(self): self.op_type = 'sub_p' X = paddle.static.data(name='X', shape=[7, 8], dtype='float64') Y = paddle.static.data(name='Y', shape=[7, 8], dtype='float64') self.input = {'X': X, 'Y': Y} self.output = { 'Z': self.layer_help.create_variable_for_type_inference(dtype=X.dtype) } self.attrs = {} self.prim2orig_args = (X, Y) self.all_ops = ['sub_p', 'elementwise_sub'] self.out_map = {self.output['Z']: 0} class TestMulPPrim2Orig(TestAddPPrim2Orig): def init_data(self): self.op_type = 'mul_p' X = paddle.static.data(name='X', shape=[7, 8], dtype='float64') Y = paddle.static.data(name='Y', shape=[7, 8], dtype='float64') self.input = {'X': X, 'Y': Y} self.output = { 'Z': self.layer_help.create_variable_for_type_inference(dtype=X.dtype) } self.attrs = {} self.prim2orig_args = (X, Y) self.all_ops = ['mul_p', 'elementwise_mul'] self.out_map = {self.output['Z']: 0} class TestDivPPrim2Orig(TestAddPPrim2Orig): def init_data(self): self.op_type = 'div_p' X = paddle.static.data(name='X', shape=[7, 8], dtype='float64') Y = paddle.static.data(name='Y', shape=[7, 8], dtype='float64') self.input = {'X': X, 'Y': Y} self.output = { 'Z': self.layer_help.create_variable_for_type_inference(dtype=X.dtype) } self.attrs = {} self.prim2orig_args = (X, Y) self.all_ops = ['div_p', 'elementwise_div'] self.out_map = {self.output['Z']: 0} class TestSqrtPPrim2Orig(TestAddPPrim2Orig): def init_data(self): self.op_type = 'sqrt_p' X = paddle.static.data(name='X', shape=[7, 8], dtype='float64') self.input = { 'X': X, } self.output = { 'Y': self.layer_help.create_variable_for_type_inference(dtype=X.dtype) } self.attrs = {} self.prim2orig_args = (X, ) self.all_ops = ['sqrt_p', 'sqrt'] self.out_map = {self.output['Y']: 0} class TestTanhPPrim2Orig(TestAddPPrim2Orig): def init_data(self): self.op_type = 'tanh_p' X = paddle.static.data(name='X', shape=[7, 8], dtype='float64') self.input = { 'X': X, } self.output = { 'Y': self.layer_help.create_variable_for_type_inference(dtype=X.dtype) } self.attrs = {} self.prim2orig_args = (X, ) self.all_ops = ['tanh_p', 'tanh'] self.out_map = {self.output['Y']: 0} class TestSinPPrim2Orig(TestAddPPrim2Orig): def init_data(self): self.op_type = 'sin_p' X = paddle.static.data(name='X', shape=[7, 8], dtype='float64') self.input = { 'X': X, } self.output = { 'Y': self.layer_help.create_variable_for_type_inference(dtype=X.dtype) } self.attrs = {} self.prim2orig_args = (X, ) self.all_ops = ['sin_p', 'sin'] self.out_map = {self.output['Y']: 0} class TestCosPPrim2Orig(TestAddPPrim2Orig): def init_data(self): self.op_type = 'cos_p' X = paddle.static.data(name='X', shape=[7, 8], dtype='float64') self.input = { 'X': X, } self.output = { 'Y': self.layer_help.create_variable_for_type_inference(dtype=X.dtype) } self.attrs = {} self.prim2orig_args = (X, ) self.all_ops = ['cos_p', 'cos'] self.out_map = {self.output['Y']: 0} class TestExpPPrim2Orig(TestAddPPrim2Orig): def init_data(self): self.op_type = 'exp_p' X = paddle.static.data(name='X', shape=[7, 8], dtype='float64') self.input = { 'X': X, } self.output = { 'Y': self.layer_help.create_variable_for_type_inference(dtype=X.dtype) } self.attrs = {} self.prim2orig_args = (X, ) self.all_ops = ['exp_p', 'exp'] self.out_map = {self.output['Y']: 0} class TestLogPPrim2Orig(TestAddPPrim2Orig): def init_data(self): self.op_type = 'log_p' X = paddle.static.data(name='X', shape=[7, 8], dtype='float64') self.input = { 'X': X, } self.output = { 'Y': self.layer_help.create_variable_for_type_inference(dtype=X.dtype) } self.attrs = {} self.prim2orig_args = (X, ) self.all_ops = ['log_p', 'log'] self.out_map = {self.output['Y']: 0} class TestReshapePPrim2Orig(TestAddPPrim2Orig): def init_data(self): self.op_type = 'reshape_p' X = paddle.static.data(name='X', shape=[2, 8], dtype='float64') self.input = { 'X': X, } self.output = { 'Y': self.layer_help.create_variable_for_type_inference(dtype=X.dtype) } self.attrs = {'shape': [4, 4]} self.prim2orig_args = (X, ) self.all_ops = ['reshape_p', 'reshape2'] self.out_map = {self.output['Y']: 0} class TestBroadcastPPrim2Orig(TestAddPPrim2Orig): def init_data(self): self.op_type = 'broadcast_p' X = paddle.static.data(name='X', shape=[2, 8], dtype='float64') self.input = { 'X': X, } self.output = { 'Y': self.layer_help.create_variable_for_type_inference(dtype=X.dtype) } self.attrs = {'shape': [10, 2, 8]} self.prim2orig_args = (X, ) self.all_ops = ['broadcast_p', 'expand_v2'] self.out_map = {self.output['Y']: 0} class TestTransposePPrim2Orig(TestAddPPrim2Orig): def init_data(self): self.op_type = 'transpose_p' X = paddle.static.data(name='X', shape=[7, 8, 9, 10], dtype='float64') self.input = { 'X': X, } self.output = { 'Y': self.layer_help.create_variable_for_type_inference(dtype=X.dtype) } self.attrs = {'axis': [1, 2, 0, 3]} self.prim2orig_args = (X, ) self.all_ops = ['transpose_p', 'transpose2'] self.out_map = {self.output['Y']: 0} class TestSplitPPrim2Orig(TestAddPPrim2Orig): def init_data(self): self.op_type = 'split_p' X = paddle.static.data(name='X', shape=[3, 9, 5], dtype='float64') self.input = { 'X': X, } self.output = { 'YS': [ self.layer_help.create_variable_for_type_inference( dtype=X.dtype) for i in range(3) ] } self.attrs = {'num_or_sections': [2, 3, 4], 'axis': 1} self.prim2orig_args = (X, ) self.all_ops = ['split_p', 'split'] self.out_map = { self.output['YS'][0]: 0, self.output['YS'][1]: 1, self.output['YS'][2]: 2, } class TestConcatPPrim2Orig(TestAddPPrim2Orig): def init_data(self): self.op_type = 'concat_p' X = paddle.static.data(name='X', shape=[3, 9, 5], dtype='float64') Y = paddle.static.data(name='Y', shape=[2, 9, 5], dtype='float64') Z = paddle.static.data(name='Z', shape=[1, 9, 5], dtype='float64') self.input = { 'XS': [X, Y, Z], } self.output = { 'Y': self.layer_help.create_variable_for_type_inference(dtype=X.dtype) } self.attrs = {'axis': 0} self.prim2orig_args = ((X, Y, Z), ) self.all_ops = ['concat_p', 'concat'] self.out_map = {self.output['Y']: 0} class TestReducePPrim2Orig(TestAddPPrim2Orig): def init_data(self): self.op_type = 'reduce_p' X = paddle.static.data(name='X', shape=[3, 9, 5], dtype='float64') self.input = {'X': X} self.output = { 'Y': self.layer_help.create_variable_for_type_inference(dtype=X.dtype) } self.attrs = {'axis': [1], 'keepdim': True} self.prim2orig_args = (X, ) self.all_ops = ['reduce_p', 'reduce_sum'] self.out_map = {self.output['Y']: 0} class TestMatmulPPrim2Orig(TestAddPPrim2Orig): def init_data(self): self.op_type = 'matmul_p' X = paddle.static.data(name='X', shape=[9, 5], dtype='float64') Y = paddle.static.data(name='Y', shape=[5, 9], dtype='float64') self.input = {'X': X, 'Y': Y} self.output = { 'Z': self.layer_help.create_variable_for_type_inference(dtype=X.dtype) } self.attrs = {} self.prim2orig_args = (X, Y) self.all_ops = ['matmul_p', 'matmul_v2'] self.out_map = {self.output['Z']: 0} class TestSliceSelectPPrim2Orig(TestAddPPrim2Orig): def init_data(self): self.op_type = 'slice_select_p' X = paddle.static.data(name='X', shape=[9, 5], dtype='float64') self.input = { 'X': X, } self.output = { 'Y': self.layer_help.create_variable_for_type_inference(dtype=X.dtype) } self.attrs = {'axis': [0], 'starts': [1], 'ends': [8], 'strides': [2]} self.prim2orig_args = (X, ) self.all_ops = ['slice_select_p', 'strided_slice'] self.out_map = {self.output['Y']: 0} class TestSliceAssignPPrim2Orig(TestAddPPrim2Orig): def init_data(self): self.op_type = 'slice_assign_p' X = paddle.static.data(name='X', shape=[9, 5], dtype='float64') Y = paddle.static.data(name='Y', shape=[9, 3], dtype='float64') self.input = {'X': X, 'Y': Y} self.output = { 'Z': self.layer_help.create_variable_for_type_inference(dtype=X.dtype) } self.attrs = {'axis': [1], 'starts': [0], 'ends': [3], 'strides': [1]} self.prim2orig_args = (X, Y) self.all_ops = ['slice_assign_p', 'assign', 'set_value'] self.out_map = {self.output['Z']: 0} class TestGatherPPrim2Orig(TestAddPPrim2Orig): def init_data(self): self.op_type = 'gather_p' X = paddle.static.data(name='X', shape=[9, 5], dtype='float64') IndexTensor = paddle.static.data(name='IndexTensor', shape=[3], dtype='int32') self.input = {'X': X, 'IndexTensor': IndexTensor} self.output = { 'Y': self.layer_help.create_variable_for_type_inference(dtype=X.dtype) } self.attrs = { 'axis': 0, } self.prim2orig_args = ( IndexTensor, X, ) self.all_ops = ['gather_p', 'gather'] self.out_map = {self.output['Y']: 0} class TestScatterAddPPrim2Orig(TestAddPPrim2Orig): def init_data(self): self.op_type = 'scatter_add_p' X = paddle.static.data(name='X', shape=[9, 5], dtype='float64') Y = paddle.static.data(name='Y', shape=[3, 5], dtype='float64') IndexTensor = paddle.static.data(name='IndexTensor', shape=[3], dtype='int32') self.input = {'X': X, 'Y': Y, 'IndexTensor': IndexTensor} self.output = { 'Z': self.layer_help.create_variable_for_type_inference(dtype=X.dtype) } self.attrs = { 'axis': 0, } self.prim2orig_args = (IndexTensor, X, Y) self.all_ops = [ 'scatter_add_p', 'fill_any_like', 'scatter', 'elementwise_add' ] self.out_map = {self.output['Z']: 0} class TestFillConstantPPrim2Orig(TestAddPPrim2Orig): def init_data(self): self.op_type = 'fill_constant_p' self.input = {} self.output = { 'Y': self.layer_help.create_variable_for_type_inference(paddle.int32) } self.attrs = {'value': 10, 'shape': [5, 5], 'dtype': paddle.int32} self.prim2orig_args = () self.all_ops = ['fill_constant_p', 'fill_constant'] self.out_map = {self.output['Y']: 0} if __name__ == '__main__': unittest.main()