mindspore 2.3.0__cp39-cp39-win_amd64.whl → 2.4.0__cp39-cp39-win_amd64.whl

mindspore/ops/auto_generate/pyboost_inner_prim.py

@@ -26,6 +26,8 @@ from mindspore._c_expression import BroadcastToPrim_
 from mindspore._c_expression import ConcatPrim_
 from mindspore._c_expression import ConvolutionGradPrim_
 from mindspore._c_expression import ConvolutionPrim_
+from mindspore._c_expression import CrossPrim_
+from mindspore._c_expression import CummaxPrim_
 from mindspore._c_expression import EluExtPrim_
 from mindspore._c_expression import FFNExtPrim_
 from mindspore._c_expression import FlashAttentionScoreGradPrim_
@@ -34,20 +36,30 @@ from mindspore._c_expression import GridSampler2DGradPrim_
 from mindspore._c_expression import GridSampler2DPrim_
 from mindspore._c_expression import GridSampler3DGradPrim_
 from mindspore._c_expression import GridSampler3DPrim_
+from mindspore._c_expression import HShrinkGradPrim_
+from mindspore._c_expression import HShrinkPrim_
+from mindspore._c_expression import IncreFlashAttentionPrim_
 from mindspore._c_expression import IsClosePrim_
+from mindspore._c_expression import LogSoftmaxGradPrim_
+from mindspore._c_expression import LogSoftmaxPrim_
 from mindspore._c_expression import MatMulPrim_
 from mindspore._c_expression import MaxPoolGradWithIndicesPrim_
 from mindspore._c_expression import MaxPoolGradWithMaskPrim_
 from mindspore._c_expression import MaxPoolWithIndicesPrim_
 from mindspore._c_expression import MaxPoolWithMaskPrim_
+from mindspore._c_expression import NanToNumPrim_
 from mindspore._c_expression import OneHotExtPrim_
 from mindspore._c_expression import ReduceAllPrim_
 from mindspore._c_expression import ReduceAnyPrim_
 from mindspore._c_expression import ReverseV2Prim_
 from mindspore._c_expression import RmsNormPrim_
+from mindspore._c_expression import RollPrim_
 from mindspore._c_expression import SearchSortedPrim_
 from mindspore._c_expression import SoftmaxPrim_
+from mindspore._c_expression import SoftShrinkGradPrim_
+from mindspore._c_expression import SoftShrinkPrim_
 from mindspore._c_expression import StackExtPrim_
+from mindspore._c_expression import TrilExtPrim_
 from mindspore._c_expression import TriuPrim_
 from mindspore._c_expression import UpsampleTrilinear3DGradPrim_
 from mindspore._c_expression import UpsampleTrilinear3DPrim_
@@ -94,8 +106,8 @@ batch_norm_grad_ext_impl = _PyboostBatchNormGradExtPrim()
 
 class _PyboostBinaryCrossEntropyGradPrim(BinaryCrossEntropyGradPrim_):
     def __call__(self, input, target, grad_output, weight, reduction):
-        converted_reduction = str_to_enum(reduction)
-        return _convert_stub(super().__call__(input, target, grad_output, weight, reduction))
+        converted_reduction = str_to_enum('binary_cross_entropy_grad', 'reduction', reduction)
+        return _convert_stub(super().__call__(input, target, grad_output, weight, converted_reduction))
 
 
 binary_cross_entropy_grad_impl = _PyboostBinaryCrossEntropyGradPrim()
@@ -103,8 +115,8 @@ binary_cross_entropy_grad_impl = _PyboostBinaryCrossEntropyGradPrim()
 
 class _PyboostBinaryCrossEntropyPrim(BinaryCrossEntropyPrim_):
     def __call__(self, input, target, weight, reduction):
-        converted_reduction = str_to_enum(reduction)
-        return _convert_stub(super().__call__(input, target, weight, reduction))
+        converted_reduction = str_to_enum('binary_cross_entropy', 'reduction', reduction)
+        return _convert_stub(super().__call__(input, target, weight, converted_reduction))
 
 
 binary_cross_entropy_impl = _PyboostBinaryCrossEntropyPrim()
@@ -112,8 +124,8 @@ binary_cross_entropy_impl = _PyboostBinaryCrossEntropyPrim()
 
 class _PyboostBCEWithLogitsLossPrim(BCEWithLogitsLossPrim_):
     def __call__(self, input, target, weight, posWeight, reduction):
-        converted_reduction = str_to_enum(reduction)
-        return _convert_stub(super().__call__(input, target, weight, posWeight, reduction))
+        converted_reduction = str_to_enum('binary_cross_entropy_with_logits', 'reduction', reduction)
+        return _convert_stub(super().__call__(input, target, weight, posWeight, converted_reduction))
 
 
 binary_cross_entropy_with_logits_impl = _PyboostBCEWithLogitsLossPrim()
@@ -139,11 +151,11 @@ concat_impl = _PyboostConcatPrim()
 
 class _PyboostConvolutionGradPrim(ConvolutionGradPrim_):
     def __call__(self, dout, input, weight, bias, stride, padding, dilation, transposed, output_padding, groups, output_mask):
-        converted_stride = to_strides(stride)
-        converted_padding = to_2d_paddings(padding)
-        converted_dilation = to_dilations(dilation)
-        converted_output_padding = to_output_padding(output_padding)
-        return _convert_stub(super().__call__(dout, input, weight, bias, stride, padding, dilation, transposed, output_padding, groups, output_mask))
+        converted_stride = to_strides('convolution_grad', 'stride', stride)
+        converted_padding = to_2d_paddings('convolution_grad', 'padding', padding)
+        converted_dilation = to_dilations('convolution_grad', 'dilation', dilation)
+        converted_output_padding = to_output_padding('convolution_grad', 'output_padding', output_padding)
+        return _convert_stub(super().__call__(dout, input, weight, bias, converted_stride, converted_padding, converted_dilation, transposed, converted_output_padding, groups, output_mask))
 
 
 convolution_grad_impl = _PyboostConvolutionGradPrim()
@@ -151,16 +163,34 @@ convolution_grad_impl = _PyboostConvolutionGradPrim()
 
 class _PyboostConvolutionPrim(ConvolutionPrim_):
     def __call__(self, input, weight, bias, stride, padding, dilation, transposed, output_padding, groups):
-        converted_stride = to_strides(stride)
-        converted_padding = to_2d_paddings(padding)
-        converted_dilation = to_dilations(dilation)
-        converted_output_padding = to_output_padding(output_padding)
-        return _convert_stub(super().__call__(input, weight, bias, stride, padding, dilation, transposed, output_padding, groups))
+        converted_stride = to_strides('convolution', 'stride', stride)
+        converted_padding = to_2d_paddings('convolution', 'padding', padding)
+        converted_dilation = to_dilations('convolution', 'dilation', dilation)
+        converted_output_padding = to_output_padding('convolution', 'output_padding', output_padding)
+        return _convert_stub(super().__call__(input, weight, bias, converted_stride, converted_padding, converted_dilation, transposed, converted_output_padding, groups))
 
 
 convolution_impl = _PyboostConvolutionPrim()
 
 
+class _PyboostCrossPrim(CrossPrim_):
+    def __call__(self, input, other, dim):
+
+        return _convert_stub(super().__call__(input, other, dim))
+
+
+cross_impl = _PyboostCrossPrim()
+
+
+class _PyboostCummaxPrim(CummaxPrim_):
+    def __call__(self, input, axis):
+
+        return _convert_stub(super().__call__(input, axis))
+
+
+cummax_impl = _PyboostCummaxPrim()
+
+
 class _PyboostEluExtPrim(EluExtPrim_):
     def __call__(self, input, alpha):
 
@@ -172,8 +202,8 @@ elu_ext_impl = _PyboostEluExtPrim()
 
 class _PyboostFFNExtPrim(FFNExtPrim_):
     def __call__(self, x, weight1, weight2, expertTokens, bias1, bias2, scale, offset, deqScale1, deqScale2, antiquant_scale1, antiquant_scale2, antiquant_offset1, antiquant_offset2, activation, inner_precise):
-        converted_activation = str_to_enum(activation)
-        return _convert_stub(super().__call__(x, weight1, weight2, expertTokens, bias1, bias2, scale, offset, deqScale1, deqScale2, antiquant_scale1, antiquant_scale2, antiquant_offset1, antiquant_offset2, activation, inner_precise))
+        converted_activation = str_to_enum('ffn_ext', 'activation', activation)
+        return _convert_stub(super().__call__(x, weight1, weight2, expertTokens, bias1, bias2, scale, offset, deqScale1, deqScale2, antiquant_scale1, antiquant_scale2, antiquant_offset1, antiquant_offset2, converted_activation, inner_precise))
 
 
 ffn_ext_impl = _PyboostFFNExtPrim()
@@ -181,8 +211,8 @@ ffn_ext_impl = _PyboostFFNExtPrim()
 
 class _PyboostFlashAttentionScoreGradPrim(FlashAttentionScoreGradPrim_):
     def __call__(self, query, key, value, dy, pse_shift, drop_mask, padding_mask, atten_mask, softmax_max, softmax_sum, softmax_in, attention_in, prefix, actual_seq_qlen, actual_seq_kvlen, head_num, keep_prob, scale_value, pre_tokens, next_tokens, inner_precise, input_layout, sparse_mode):
-        converted_input_layout = str_to_enum(input_layout)
-        return _convert_stub(super().__call__(query, key, value, dy, pse_shift, drop_mask, padding_mask, atten_mask, softmax_max, softmax_sum, softmax_in, attention_in, prefix, actual_seq_qlen, actual_seq_kvlen, head_num, keep_prob, scale_value, pre_tokens, next_tokens, inner_precise, input_layout, sparse_mode))
+        converted_input_layout = str_to_enum('flash_attention_score_grad', 'input_layout', input_layout)
+        return _convert_stub(super().__call__(query, key, value, dy, pse_shift, drop_mask, padding_mask, atten_mask, softmax_max, softmax_sum, softmax_in, attention_in, prefix, actual_seq_qlen, actual_seq_kvlen, head_num, keep_prob, scale_value, pre_tokens, next_tokens, inner_precise, converted_input_layout, sparse_mode))
 
 
 flash_attention_score_grad_impl = _PyboostFlashAttentionScoreGradPrim()
@@ -190,8 +220,8 @@ flash_attention_score_grad_impl = _PyboostFlashAttentionScoreGradPrim()
 
 class _PyboostFlashAttentionScorePrim(FlashAttentionScorePrim_):
     def __call__(self, query, key, value, real_shift, drop_mask, padding_mask, attn_mask, prefix, actual_seq_qlen, actual_seq_kvlen, head_num, keep_prob, scale_value, pre_tokens, next_tokens, inner_precise, input_layout, sparse_mode):
-        converted_input_layout = str_to_enum(input_layout)
-        return _convert_stub(super().__call__(query, key, value, real_shift, drop_mask, padding_mask, attn_mask, prefix, actual_seq_qlen, actual_seq_kvlen, head_num, keep_prob, scale_value, pre_tokens, next_tokens, inner_precise, input_layout, sparse_mode))
+        converted_input_layout = str_to_enum('flash_attention_score', 'input_layout', input_layout)
+        return _convert_stub(super().__call__(query, key, value, real_shift, drop_mask, padding_mask, attn_mask, prefix, actual_seq_qlen, actual_seq_kvlen, head_num, keep_prob, scale_value, pre_tokens, next_tokens, inner_precise, converted_input_layout, sparse_mode))
 
 
 flash_attention_score_impl = _PyboostFlashAttentionScorePrim()
@@ -199,9 +229,9 @@ flash_attention_score_impl = _PyboostFlashAttentionScorePrim()
 
 class _PyboostGridSampler2DGradPrim(GridSampler2DGradPrim_):
     def __call__(self, grad, input_x, grid, interpolation_mode, padding_mode, align_corners):
-        converted_interpolation_mode = str_to_enum(interpolation_mode)
-        converted_padding_mode = str_to_enum(padding_mode)
-        return _convert_stub(super().__call__(grad, input_x, grid, interpolation_mode, padding_mode, align_corners))
+        converted_interpolation_mode = str_to_enum('grid_sampler_2d_grad', 'interpolation_mode', interpolation_mode)
+        converted_padding_mode = str_to_enum('grid_sampler_2d_grad', 'padding_mode', padding_mode)
+        return _convert_stub(super().__call__(grad, input_x, grid, converted_interpolation_mode, converted_padding_mode, align_corners))
 
 
 grid_sampler_2d_grad_impl = _PyboostGridSampler2DGradPrim()
@@ -209,9 +239,9 @@ grid_sampler_2d_grad_impl = _PyboostGridSampler2DGradPrim()
 
 class _PyboostGridSampler2DPrim(GridSampler2DPrim_):
     def __call__(self, input_x, grid, interpolation_mode, padding_mode, align_corners):
-        converted_interpolation_mode = str_to_enum(interpolation_mode)
-        converted_padding_mode = str_to_enum(padding_mode)
-        return _convert_stub(super().__call__(input_x, grid, interpolation_mode, padding_mode, align_corners))
+        converted_interpolation_mode = str_to_enum('grid_sampler_2d', 'interpolation_mode', interpolation_mode)
+        converted_padding_mode = str_to_enum('grid_sampler_2d', 'padding_mode', padding_mode)
+        return _convert_stub(super().__call__(input_x, grid, converted_interpolation_mode, converted_padding_mode, align_corners))
 
 
 grid_sampler_2d_impl = _PyboostGridSampler2DPrim()
@@ -219,9 +249,9 @@ grid_sampler_2d_impl = _PyboostGridSampler2DPrim()
 
 class _PyboostGridSampler3DGradPrim(GridSampler3DGradPrim_):
     def __call__(self, grad, input_x, grid, interpolation_mode, padding_mode, align_corners):
-        converted_interpolation_mode = str_to_enum(interpolation_mode)
-        converted_padding_mode = str_to_enum(padding_mode)
-        return _convert_stub(super().__call__(grad, input_x, grid, interpolation_mode, padding_mode, align_corners))
+        converted_interpolation_mode = str_to_enum('grid_sampler_3d_grad', 'interpolation_mode', interpolation_mode)
+        converted_padding_mode = str_to_enum('grid_sampler_3d_grad', 'padding_mode', padding_mode)
+        return _convert_stub(super().__call__(grad, input_x, grid, converted_interpolation_mode, converted_padding_mode, align_corners))
 
 
 grid_sampler_3d_grad_impl = _PyboostGridSampler3DGradPrim()
@@ -229,14 +259,41 @@ grid_sampler_3d_grad_impl = _PyboostGridSampler3DGradPrim()
 
 class _PyboostGridSampler3DPrim(GridSampler3DPrim_):
     def __call__(self, input_x, grid, interpolation_mode, padding_mode, align_corners):
-        converted_interpolation_mode = str_to_enum(interpolation_mode)
-        converted_padding_mode = str_to_enum(padding_mode)
-        return _convert_stub(super().__call__(input_x, grid, interpolation_mode, padding_mode, align_corners))
+        converted_interpolation_mode = str_to_enum('grid_sampler_3d', 'interpolation_mode', interpolation_mode)
+        converted_padding_mode = str_to_enum('grid_sampler_3d', 'padding_mode', padding_mode)
+        return _convert_stub(super().__call__(input_x, grid, converted_interpolation_mode, converted_padding_mode, align_corners))
 
 
 grid_sampler_3d_impl = _PyboostGridSampler3DPrim()
 
 
+class _PyboostHShrinkGradPrim(HShrinkGradPrim_):
+    def __call__(self, gradients, features, lambd):
+
+        return _convert_stub(super().__call__(gradients, features, lambd))
+
+
+hshrink_grad_impl = _PyboostHShrinkGradPrim()
+
+
+class _PyboostHShrinkPrim(HShrinkPrim_):
+    def __call__(self, input, lambd):
+
+        return _convert_stub(super().__call__(input, lambd))
+
+
+hshrink_impl = _PyboostHShrinkPrim()
+
+
+class _PyboostIncreFlashAttentionPrim(IncreFlashAttentionPrim_):
+    def __call__(self, query, key, value, attn_mask, actual_seq_lengths, pse_shift, dequant_scale1, quant_scale1, dequant_scale2, quant_scale2, quant_offset2, antiquant_scale, antiquant_offset, block_table, kv_padding_size, num_heads, input_layout, scale_value, num_key_value_heads, block_size, inner_precise):
+        converted_input_layout = str_to_enum('incre_flash_attention', 'input_layout', input_layout)
+        return _convert_stub(super().__call__(query, key, value, attn_mask, actual_seq_lengths, pse_shift, dequant_scale1, quant_scale1, dequant_scale2, quant_scale2, quant_offset2, antiquant_scale, antiquant_offset, block_table, kv_padding_size, num_heads, converted_input_layout, scale_value, num_key_value_heads, block_size, inner_precise))
+
+
+incre_flash_attention_impl = _PyboostIncreFlashAttentionPrim()
+
+
 class _PyboostIsClosePrim(IsClosePrim_):
     def __call__(self, input, other, rtol, atol, equal_nan):
 
@@ -246,6 +303,24 @@ class _PyboostIsClosePrim(IsClosePrim_):
 isclose_impl = _PyboostIsClosePrim()
 
 
+class _PyboostLogSoftmaxGradPrim(LogSoftmaxGradPrim_):
+    def __call__(self, logits, grad, axis):
+
+        return _convert_stub(super().__call__(logits, grad, axis))
+
+
+log_softmax_grad_impl = _PyboostLogSoftmaxGradPrim()
+
+
+class _PyboostLogSoftmaxPrim(LogSoftmaxPrim_):
+    def __call__(self, logits, axis):
+
+        return _convert_stub(super().__call__(logits, axis))
+
+
+log_softmax_impl = _PyboostLogSoftmaxPrim()
+
+
 class _PyboostMatMulPrim(MatMulPrim_):
     def __call__(self, input, mat2, transpose_a, transpose_b):
 
@@ -257,11 +332,11 @@ matmul_impl = _PyboostMatMulPrim()
 
 class _PyboostMaxPoolGradWithIndicesPrim(MaxPoolGradWithIndicesPrim_):
     def __call__(self, x, grad, argmax, kernel_size, strides, pads, dilation, ceil_mode, argmax_type):
-        converted_kernel_size = to_kernel_size(kernel_size)
-        converted_strides = to_strides(strides)
-        converted_pads = to_output_padding(pads)
-        converted_dilation = to_dilations(dilation)
-        return _convert_stub(super().__call__(x, grad, argmax, kernel_size, strides, pads, dilation, ceil_mode, argmax_type))
+        converted_kernel_size = to_kernel_size('max_pool_grad_with_indices', 'kernel_size', kernel_size)
+        converted_strides = to_strides('max_pool_grad_with_indices', 'strides', strides)
+        converted_pads = to_output_padding('max_pool_grad_with_indices', 'pads', pads)
+        converted_dilation = to_dilations('max_pool_grad_with_indices', 'dilation', dilation)
+        return _convert_stub(super().__call__(x, grad, argmax, converted_kernel_size, converted_strides, converted_pads, converted_dilation, ceil_mode, argmax_type))
 
 
 max_pool_grad_with_indices_impl = _PyboostMaxPoolGradWithIndicesPrim()
@@ -269,11 +344,11 @@ max_pool_grad_with_indices_impl = _PyboostMaxPoolGradWithIndicesPrim()
 
 class _PyboostMaxPoolGradWithMaskPrim(MaxPoolGradWithMaskPrim_):
     def __call__(self, x, grad, mask, kernel_size, strides, pads, dilation, ceil_mode, argmax_type):
-        converted_kernel_size = to_kernel_size(kernel_size)
-        converted_strides = to_strides(strides)
-        converted_pads = to_output_padding(pads)
-        converted_dilation = to_dilations(dilation)
-        return _convert_stub(super().__call__(x, grad, mask, kernel_size, strides, pads, dilation, ceil_mode, argmax_type))
+        converted_kernel_size = to_kernel_size('max_pool_grad_with_mask', 'kernel_size', kernel_size)
+        converted_strides = to_strides('max_pool_grad_with_mask', 'strides', strides)
+        converted_pads = to_output_padding('max_pool_grad_with_mask', 'pads', pads)
+        converted_dilation = to_dilations('max_pool_grad_with_mask', 'dilation', dilation)
+        return _convert_stub(super().__call__(x, grad, mask, converted_kernel_size, converted_strides, converted_pads, converted_dilation, ceil_mode, argmax_type))
 
 
 max_pool_grad_with_mask_impl = _PyboostMaxPoolGradWithMaskPrim()
@@ -281,11 +356,11 @@ max_pool_grad_with_mask_impl = _PyboostMaxPoolGradWithMaskPrim()
 
 class _PyboostMaxPoolWithIndicesPrim(MaxPoolWithIndicesPrim_):
     def __call__(self, x, kernel_size, strides, pads, dilation, ceil_mode, argmax_type):
-        converted_kernel_size = to_kernel_size(kernel_size)
-        converted_strides = to_strides(strides)
-        converted_pads = to_output_padding(pads)
-        converted_dilation = to_dilations(dilation)
-        return _convert_stub(super().__call__(x, kernel_size, strides, pads, dilation, ceil_mode, argmax_type))
+        converted_kernel_size = to_kernel_size('max_pool_with_indices', 'kernel_size', kernel_size)
+        converted_strides = to_strides('max_pool_with_indices', 'strides', strides)
+        converted_pads = to_output_padding('max_pool_with_indices', 'pads', pads)
+        converted_dilation = to_dilations('max_pool_with_indices', 'dilation', dilation)
+        return _convert_stub(super().__call__(x, converted_kernel_size, converted_strides, converted_pads, converted_dilation, ceil_mode, argmax_type))
 
 
 max_pool_with_indices_impl = _PyboostMaxPoolWithIndicesPrim()
@@ -293,16 +368,25 @@ max_pool_with_indices_impl = _PyboostMaxPoolWithIndicesPrim()
 
 class _PyboostMaxPoolWithMaskPrim(MaxPoolWithMaskPrim_):
     def __call__(self, x, kernel_size, strides, pads, dilation, ceil_mode, argmax_type):
-        converted_kernel_size = to_kernel_size(kernel_size)
-        converted_strides = to_strides(strides)
-        converted_pads = to_output_padding(pads)
-        converted_dilation = to_dilations(dilation)
-        return _convert_stub(super().__call__(x, kernel_size, strides, pads, dilation, ceil_mode, argmax_type))
+        converted_kernel_size = to_kernel_size('max_pool_with_mask', 'kernel_size', kernel_size)
+        converted_strides = to_strides('max_pool_with_mask', 'strides', strides)
+        converted_pads = to_output_padding('max_pool_with_mask', 'pads', pads)
+        converted_dilation = to_dilations('max_pool_with_mask', 'dilation', dilation)
+        return _convert_stub(super().__call__(x, converted_kernel_size, converted_strides, converted_pads, converted_dilation, ceil_mode, argmax_type))
 
 
 max_pool_with_mask_impl = _PyboostMaxPoolWithMaskPrim()
 
 
+class _PyboostNanToNumPrim(NanToNumPrim_):
+    def __call__(self, input, nan, posinf, neginf):
+
+        return _convert_stub(super().__call__(input, nan, posinf, neginf))
+
+
+nan_to_num_impl = _PyboostNanToNumPrim()
+
+
 class _PyboostOneHotExtPrim(OneHotExtPrim_):
     def __call__(self, tensor, num_classes, on_value, off_value, axis):
 
@@ -348,6 +432,15 @@ class _PyboostRmsNormPrim(RmsNormPrim_):
 rms_norm_impl = _PyboostRmsNormPrim()
 
 
+class _PyboostRollPrim(RollPrim_):
+    def __call__(self, input, shift, axis):
+
+        return _convert_stub(super().__call__(input, shift, axis))
+
+
+roll_impl = _PyboostRollPrim()
+
+
 class _PyboostSearchSortedPrim(SearchSortedPrim_):
     def __call__(self, sorted_sequence, values, sorter, dtype, right):
 
@@ -366,6 +459,24 @@ class _PyboostSoftmaxPrim(SoftmaxPrim_):
 softmax_impl = _PyboostSoftmaxPrim()
 
 
+class _PyboostSoftShrinkGradPrim(SoftShrinkGradPrim_):
+    def __call__(self, input_grad, input_x, lambd):
+
+        return _convert_stub(super().__call__(input_grad, input_x, lambd))
+
+
+softshrink_grad_impl = _PyboostSoftShrinkGradPrim()
+
+
+class _PyboostSoftShrinkPrim(SoftShrinkPrim_):
+    def __call__(self, input, lambd):
+
+        return _convert_stub(super().__call__(input, lambd))
+
+
+softshrink_impl = _PyboostSoftShrinkPrim()
+
+
 class _PyboostStackExtPrim(StackExtPrim_):
     def __call__(self, tensors, dim):
 
@@ -375,6 +486,15 @@ class _PyboostStackExtPrim(StackExtPrim_):
 stack_ext_impl = _PyboostStackExtPrim()
 
 
+class _PyboostTrilExtPrim(TrilExtPrim_):
+    def __call__(self, input, diagonal):
+
+        return _convert_stub(super().__call__(input, diagonal))
+
+
+tril_ext_impl = _PyboostTrilExtPrim()
+
+
 class _PyboostTriuPrim(TriuPrim_):
     def __call__(self, input, diagonal):
 
@@ -412,9 +532,9 @@ grouped_matmul_impl = _PyboostGroupedMatmulPrim()
 
 
 class _PyboostQuantBatchMatmulPrim(QuantBatchMatmulPrim_):
-    def __call__(self, x1, x2, scale, offset, bias, transpose_x1, transpose_x2, dtype):
+    def __call__(self, x1, x2, scale, offset, bias, pertokenScaleOptional, transpose_x1, transpose_x2, dtype):
 
-        return _convert_stub(super().__call__(x1, x2, scale, offset, bias, transpose_x1, transpose_x2, dtype))
+        return _convert_stub(super().__call__(x1, x2, scale, offset, bias, pertokenScaleOptional, transpose_x1, transpose_x2, dtype))
 
 
 quant_batch_matmul_impl = _PyboostQuantBatchMatmulPrim()

mindspore/ops/composite/base.py

@@ -30,7 +30,7 @@ from mindspore._c_expression import GradOperation_, HyperMap_, Map_, MultitypeFu
     SequenceSliceGetItem_, ListSliceSetItem_, VmapOperation_, TaylorOperation_, ListPop_, \
     ListClear_, ListReverse_, ListExtend_, DictClear_, DictHasKey_, DictUpdate_, DictFromKeys_, \
     ZerosLike_, TensorIndexGetitem_, TensorIndexSetitem_, ListAdd_, DictSetItem_, \
-    HandleBoolTensor_, PreSetitemByTuple_, StarredGetItem_,\
+    HandleBoolTensor_, PreSetitemByTuple_, StarredGetItem_, \
     StarredUnpack_, StarredUnpackMerge_, IterConverter_, HasNext_, Next_, MSContext
 from mindspore.common import dtype as mstype
 from mindspore.common.api import jit, _pynative_executor, _wrap_func
@@ -346,9 +346,11 @@ class GradOperation(GradOperation_):
         self.grad_position = (0,)
 
     def __call__(self, fn, weights=None):
-        weights_id = _get_grad_weights_id(weights)
-        if self.grad_fn is not None and self.fn == fn and self.weights_id == weights_id:
-            return self.grad_fn
+        weights_id = ''
+        if context.get_context("mode") == context.GRAPH_MODE:
+            weights_id = _get_grad_weights_id(weights)
+            if self.grad_fn is not None and self.fn == fn and self.weights_id == weights_id:
+                return self.grad_fn
         grad_ = GradOperation(self.get_all, self.get_by_list, self.sens_param)
         # If calling Grad in GRAPH_MODE or calling Grad in functions decorated with 'jit', do grad in GRAPH_MODE
         # If calling Grad in pure PYNATIVE_MODE do grad in PYNATIVE_MODE
@@ -374,8 +376,8 @@ class GradOperation(GradOperation_):
 
             @_wrap_func
             def after_grad(*args, **kwargs):
-                self._pynative_forward_run(fn, grad_, weights, args, kwargs)
-                out = _pynative_executor.grad(fn, grad_, weights, self.grad_position, *args, **kwargs)
+                run_args = self._pynative_forward_run(fn, grad_, weights, *args, **kwargs)
+                out = _pynative_executor.grad(fn, grad_, weights, self.grad_position, *run_args)
                 out = _grads_divided_by_device_num_if_recomputation(out)
                 return out
         else:
@@ -396,26 +398,39 @@ class GradOperation(GradOperation_):
         self.weights_id = weights_id
         return self.grad_fn
 
-    def _pynative_forward_run(self, fn, grad, weights, args, kwargs):
-        """ Pynative forward run to build grad graph. """
-        new_kwargs = kwargs
+    def _pynative_forward_run(self, fn, grad, weights, *args, **kwargs):
+        """ PyNative forward run to build grad graph. """
+        sens = None
         if self.sens_param:
-            if 'sens' not in kwargs.keys():
-                args = args[:-1]
+            if 'sens' in kwargs.keys():
+                sens = kwargs.pop('sens')
             else:
-                new_kwargs = kwargs.copy()
-                new_kwargs.pop('sens')
+                # default use args last elem as sens
+                sens = args[-1]
+                args = args[:-1]
+        run_args = args
+        if kwargs:
+            run_args = args + tuple(kwargs.values())
+
+        # check run exclude sens
         if isinstance(fn, (FunctionType, MethodType)):
-            if not _pynative_executor.check_run(grad, fn, weights, None, *args, **new_kwargs):
+            if not _pynative_executor.check_run(grad, fn, weights, None, *run_args):
                 _pynative_executor.set_grad_flag(True)
-                _pynative_executor.new_graph(fn, *args, **new_kwargs)
-                output = fn(*args, **new_kwargs)
-                _pynative_executor.end_graph(fn, output, *args, **new_kwargs)
+                _pynative_executor.new_graph(fn, *args, **kwargs)
+                output = fn(*args, **kwargs)
+                _pynative_executor.end_graph(fn, output, *args, **kwargs)
         else:
-            # Check if fn have run already
-            if not _pynative_executor.check_run(grad, fn, weights, None, *args, **new_kwargs):
-                _pynative_executor.set_grad_flag(True)
-                fn(*args, **new_kwargs)
+            # Check if fn has run already
+            if not _pynative_executor.check_run(grad, fn, weights, None, *run_args):
+                requires_grad = fn.requires_grad
+                fn.requires_grad = True
+                fn(*args, **kwargs)
+                fn.requires_grad = requires_grad
+
+        # If it has sens, keep sens as the last element
+        if sens is not None:
+            run_args += (sens,) if sens is not isinstance(run_args, tuple) else sens
+        return run_args
 
 
 class _TaylorOperation(TaylorOperation_):
@@ -552,13 +567,15 @@ class _Grad(GradOperation_):
         self.weights_id = None
 
     def __call__(self, fn, weights=None, grad_position=0):
-        weights_id = _get_grad_weights_id(weights)
-        if self.grad_fn is not None and self.fn == fn and self.grad_position == grad_position and \
-                self.weights_id == weights_id:
-            return self.grad_fn
+        weights_id = ''
+        if context.get_context("mode") == context.GRAPH_MODE:
+            weights_id = _get_grad_weights_id(weights)
+            if self.grad_fn is not None and self.fn == fn and self.grad_position == grad_position and \
+                    self.weights_id == weights_id:
+                return self.grad_fn
 
-        def aux_fn(*args):
-            outputs = fn(*args)
+        def aux_fn(*args, **kwargs):
+            outputs = fn(*args, **kwargs)
             if not isinstance(outputs, tuple) or len(outputs) < 2:
                 raise ValueError("When has_aux is True, origin fn requires more than one outputs.")
             res = (outputs[0],)
@@ -597,8 +614,8 @@ class _Grad(GradOperation_):
 
             @_wrap_func
             def after_grad(*args, **kwargs):
-                res = self._pynative_forward_run(fn, grad_, weights, args, kwargs)
-                out = _pynative_executor.grad(fn, grad_, weights, grad_position, *args, **kwargs)
+                run_args, res = self._pynative_forward_run(fn, grad_, weights, *args, **kwargs)
+                out = _pynative_executor.grad(fn, grad_, weights, grad_position, *run_args)
                 out = _grads_divided_by_device_num_if_recomputation(out)
                 if self.return_ids and out:
                     out = _combine_with_ids(grad_position, weights, out)
@@ -633,32 +650,49 @@ class _Grad(GradOperation_):
         self.weights_id = weights_id
         return self.grad_fn
 
-    def _pynative_forward_run(self, fn, grad, weights, args, kwargs):
-        """ Pynative forward runs to build grad graph. """
-        new_kwargs = kwargs
-        outputs = ()
+    def _pynative_forward_run(self, fn, grad, weights, *args, **kwargs):
+        """ PyNative forward runs to build grad graph. """
+        sens = None
         if self.sens_param:
             if 'sens' in kwargs.keys():
-                new_kwargs = kwargs.copy()
-                new_kwargs.pop('sens')
+                sens = kwargs.pop('sens')
             else:
+                # default use args last elem as sens
+                sens = args[-1]
                 args = args[:-1]
+        run_args = args
+        if kwargs:
+            run_args = args + tuple(kwargs.values())
+
+        # check run exclude sens
+        outputs = ()
+        run_forward = False
         if isinstance(fn, (FunctionType, MethodType)):
-            if not _pynative_executor.check_run(grad, fn, weights, self.grad_position, *args, **new_kwargs):
+            if not _pynative_executor.check_run(grad, fn, weights, self.grad_position, *run_args):
                 _pynative_executor.set_grad_flag(True)
-                _pynative_executor.new_graph(fn, *args, **new_kwargs)
-                outputs = fn(*args, **new_kwargs)
-                _pynative_executor.end_graph(fn, outputs, *args, **new_kwargs)
-                return outputs
+                _pynative_executor.new_graph(fn, *args, **kwargs)
+                outputs = fn(*args, **kwargs)
+                _pynative_executor.end_graph(fn, outputs, *args, **kwargs)
+                run_forward = True
         else:
             # Check if fn has run already.
-            if not _pynative_executor.check_run(grad, fn, weights, self.grad_position, *args, **new_kwargs):
-                _pynative_executor.set_grad_flag(True)
-                outputs = fn(*args, **new_kwargs)
-                return outputs
+            if not _pynative_executor.check_run(grad, fn, weights, self.grad_position, *run_args):
+                requires_grad = fn.requires_grad
+                fn.requires_grad = True
+                outputs = fn(*args, **kwargs)
+                fn.requires_grad = requires_grad
+                run_forward = True
+        # If it has sens, keep sens as the last element
+        if sens is not None:
+            run_args += (sens,) if sens is not isinstance(run_args, tuple) else sens
+
+        # Normal run grad
+        if run_forward:
+            return run_args, outputs
+
         if (self.get_value or self.has_aux) and not outputs:
-            outputs = fn(*args, **new_kwargs)
-        return outputs
+            outputs = fn(*args, **kwargs)
+        return run_args, outputs
 
 
 class _Vmap(VmapOperation_):
@@ -806,10 +840,12 @@ class MultitypeFuncGraph(MultitypeFuncGraph_):
 
 class HyperMap(HyperMap_):
     """
-    Hypermap will apply the set operation to input sequences.
+    HyperMap will apply the set operation to input sequences.
 
     Apply the operations to every element of the sequence or nested sequence. Different
-    from `mindspore.ops.Map`, the `HyperMap` supports to apply on nested structure.
+    from `mindspore.ops.Map`, the `HyperMap` supports to apply on nested structure. The
+    `HyperMap` also supports dynamic sequences as input, but it does not extend this
+    support to nested dynamic sequences.
 
     Args:
         ops (Union[MultitypeFuncGraph, None], optional): `ops` is the operation to apply. If `ops` is `None`,
@@ -959,6 +995,7 @@ class _ListAppend(ListAppend_):
     Args:
         name (str): The name of the metafuncgraph object.
     """
+
     # `__init__` method removed entirely
     def __call__(self, *args):
         pass

mindspore/ops/composite/multitype_ops/_compile_utils.py

@@ -483,6 +483,7 @@ def format_index_tensor(index, arg):
             index[format_idx] = F.select(index_tensor < 0, index_tensor + format_dim, index_tensor)
         return index
     index = Tensor(index)
+    format_dims = Tensor(format_dims)
     return F.select(index < 0, index + format_dims, index)
 
 

mindspore/ops/composite/multitype_ops/not_in_impl.py

@@ -41,7 +41,7 @@ def _number_not_in_tuple(x, y):
     Returns:
        bool, if x not in y return true, x in y return false.
    """
-    if F.is_sequence_shape_unknown(y) or not F.isconstant(x):
+    if F.is_sequence_value_unknown(y) or not F.isconstant(x):
         return not InSequence()(x, y)
     return not const_utils.scalar_in_sequence(x, y)
 
@@ -58,7 +58,7 @@ def _number_not_in_list(x, y):
     Returns:
        bool, if x not in y return true, x in y return false.
    """
-    if F.is_sequence_shape_unknown(y) or not F.isconstant(x):
+    if F.is_sequence_value_unknown(y) or not F.isconstant(x):
         return not InSequence()(x, y)
     return not const_utils.scalar_in_sequence(x, y)