mindspore 2.4.0__cp311-cp311-macosx_10_15_x86_64.whl

mindspore/ops/auto_generate/pyboost_inner_prim.py

@@ -0,0 +1,549 @@
+# Copyright 2023 Huawei Technologies Co., Ltd
+#
+# 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.
+# ============================================================================
+
+from mindspore.common._stub_tensor import _convert_stub
+from mindspore.ops.auto_generate.gen_arg_handler import *
+from mindspore._c_expression import ArgMaxWithValuePrim_
+from mindspore._c_expression import ArgMinWithValuePrim_
+from mindspore._c_expression import BatchMatMulPrim_
+from mindspore._c_expression import BatchNormGradExtPrim_
+from mindspore._c_expression import BinaryCrossEntropyGradPrim_
+from mindspore._c_expression import BinaryCrossEntropyPrim_
+from mindspore._c_expression import BCEWithLogitsLossPrim_
+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_
+from mindspore._c_expression import FlashAttentionScorePrim_
+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_
+from mindspore._c_expression import GroupedMatmulPrim_
+from mindspore._c_expression import QuantBatchMatmulPrim_
+from mindspore._c_expression import WeightQuantBatchMatmulPrim_
+
+
+class _PyboostArgMaxWithValuePrim(ArgMaxWithValuePrim_):
+    def __call__(self, input, axis, keep_dims):
+
+        return _convert_stub(super().__call__(input, axis, keep_dims))
+
+
+argmax_with_value_impl = _PyboostArgMaxWithValuePrim()
+
+
+class _PyboostArgMinWithValuePrim(ArgMinWithValuePrim_):
+    def __call__(self, input, axis, keep_dims):
+
+        return _convert_stub(super().__call__(input, axis, keep_dims))
+
+
+argmin_with_value_impl = _PyboostArgMinWithValuePrim()
+
+
+class _PyboostBatchMatMulPrim(BatchMatMulPrim_):
+    def __call__(self, x, y, transpose_a, transpose_b):
+
+        return _convert_stub(super().__call__(x, y, transpose_a, transpose_b))
+
+
+batch_mat_mul_impl = _PyboostBatchMatMulPrim()
+
+
+class _PyboostBatchNormGradExtPrim(BatchNormGradExtPrim_):
+    def __call__(self, dout, input, weight, running_mean, running_var, saved_mean, saved_rstd, training, eps):
+
+        return _convert_stub(super().__call__(dout, input, weight, running_mean, running_var, saved_mean, saved_rstd, training, eps))
+
+
+batch_norm_grad_ext_impl = _PyboostBatchNormGradExtPrim()
+
+
+class _PyboostBinaryCrossEntropyGradPrim(BinaryCrossEntropyGradPrim_):
+    def __call__(self, 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()
+
+
+class _PyboostBinaryCrossEntropyPrim(BinaryCrossEntropyPrim_):
+    def __call__(self, 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()
+
+
+class _PyboostBCEWithLogitsLossPrim(BCEWithLogitsLossPrim_):
+    def __call__(self, 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()
+
+
+class _PyboostBroadcastToPrim(BroadcastToPrim_):
+    def __call__(self, input, shape):
+
+        return _convert_stub(super().__call__(input, shape))
+
+
+broadcast_to_impl = _PyboostBroadcastToPrim()
+
+
+class _PyboostConcatPrim(ConcatPrim_):
+    def __call__(self, tensors, axis):
+
+        return _convert_stub(super().__call__(tensors, axis))
+
+
+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('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()
+
+
+class _PyboostConvolutionPrim(ConvolutionPrim_):
+    def __call__(self, 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):
+
+        return _convert_stub(super().__call__(input, alpha))
+
+
+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('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()
+
+
+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('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()
+
+
+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('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()
+
+
+class _PyboostGridSampler2DGradPrim(GridSampler2DGradPrim_):
+    def __call__(self, 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()
+
+
+class _PyboostGridSampler2DPrim(GridSampler2DPrim_):
+    def __call__(self, 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()
+
+
+class _PyboostGridSampler3DGradPrim(GridSampler3DGradPrim_):
+    def __call__(self, 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()
+
+
+class _PyboostGridSampler3DPrim(GridSampler3DPrim_):
+    def __call__(self, 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):
+
+        return _convert_stub(super().__call__(input, other, rtol, atol, equal_nan))
+
+
+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):
+
+        return _convert_stub(super().__call__(input, mat2, transpose_a, transpose_b))
+
+
+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('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()
+
+
+class _PyboostMaxPoolGradWithMaskPrim(MaxPoolGradWithMaskPrim_):
+    def __call__(self, 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()
+
+
+class _PyboostMaxPoolWithIndicesPrim(MaxPoolWithIndicesPrim_):
+    def __call__(self, 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()
+
+
+class _PyboostMaxPoolWithMaskPrim(MaxPoolWithMaskPrim_):
+    def __call__(self, 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):
+
+        return _convert_stub(super().__call__(tensor, num_classes, on_value, off_value, axis))
+
+
+one_hot_ext_impl = _PyboostOneHotExtPrim()
+
+
+class _PyboostReduceAllPrim(ReduceAllPrim_):
+    def __call__(self, input, axis, keep_dims):
+
+        return _convert_stub(super().__call__(input, axis, keep_dims))
+
+
+reduce_all_impl = _PyboostReduceAllPrim()
+
+
+class _PyboostReduceAnyPrim(ReduceAnyPrim_):
+    def __call__(self, x, axis, keep_dims):
+
+        return _convert_stub(super().__call__(x, axis, keep_dims))
+
+
+reduce_any_impl = _PyboostReduceAnyPrim()
+
+
+class _PyboostReverseV2Prim(ReverseV2Prim_):
+    def __call__(self, input, axis):
+
+        return _convert_stub(super().__call__(input, axis))
+
+
+reverse_v2_impl = _PyboostReverseV2Prim()
+
+
+class _PyboostRmsNormPrim(RmsNormPrim_):
+    def __call__(self, x, gamma, epsilon):
+
+        return _convert_stub(super().__call__(x, gamma, epsilon))
+
+
+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):
+
+        return _convert_stub(super().__call__(sorted_sequence, values, sorter, dtype, right))
+
+
+searchsorted_impl = _PyboostSearchSortedPrim()
+
+
+class _PyboostSoftmaxPrim(SoftmaxPrim_):
+    def __call__(self, input, axis):
+
+        return _convert_stub(super().__call__(input, axis))
+
+
+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):
+
+        return _convert_stub(super().__call__(tensors, dim))
+
+
+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):
+
+        return _convert_stub(super().__call__(input, diagonal))
+
+
+triu_impl = _PyboostTriuPrim()
+
+
+class _PyboostUpsampleTrilinear3DGradPrim(UpsampleTrilinear3DGradPrim_):
+    def __call__(self, dy, input_size, output_size, scales, align_corners):
+
+        return _convert_stub(super().__call__(dy, input_size, output_size, scales, align_corners))
+
+
+upsample_trilinear3d_grad_impl = _PyboostUpsampleTrilinear3DGradPrim()
+
+
+class _PyboostUpsampleTrilinear3DPrim(UpsampleTrilinear3DPrim_):
+    def __call__(self, x, output_size, scales, align_corners):
+
+        return _convert_stub(super().__call__(x, output_size, scales, align_corners))
+
+
+upsample_trilinear3d_impl = _PyboostUpsampleTrilinear3DPrim()
+
+
+class _PyboostGroupedMatmulPrim(GroupedMatmulPrim_):
+    def __call__(self, x, weight, bias, scale, offset, antiquant_scale, antiquant_offset, group_list, split_item, group_type):
+
+        return _convert_stub(super().__call__(x, weight, bias, scale, offset, antiquant_scale, antiquant_offset, group_list, split_item, group_type))
+
+
+grouped_matmul_impl = _PyboostGroupedMatmulPrim()
+
+
+class _PyboostQuantBatchMatmulPrim(QuantBatchMatmulPrim_):
+    def __call__(self, x1, x2, scale, offset, bias, pertokenScaleOptional, 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()
+
+
+class _PyboostWeightQuantBatchMatmulPrim(WeightQuantBatchMatmulPrim_):
+    def __call__(self, x, weight, antiquant_scale, antiquant_offset, quant_scale, quant_offset, bias, transpose_x, transpose_weight, antiquant_group_size):
+
+        return _convert_stub(super().__call__(x, weight, antiquant_scale, antiquant_offset, quant_scale, quant_offset, bias, transpose_x, transpose_weight, antiquant_group_size))
+
+
+weight_quant_batch_matmul_impl = _PyboostWeightQuantBatchMatmulPrim()

mindspore/ops/composite/__init__.py

@@ -0,0 +1,71 @@
+# Copyright 2020 Huawei Technologies Co., Ltd
+#
+# 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.
+# ============================================================================
+
+"""
+Composite operators.
+
+Pre-defined combination of operators.
+"""
+
+from __future__ import absolute_import
+from mindspore.ops.composite.base import GradOperation, _Grad, HyperMap, Map, MultitypeFuncGraph, add_flags, \
+    tail, zip_operation, _Vmap, _TaylorOperation, iter_converter, ms_hasnext, ms_next
+from mindspore.ops.composite.env_ops import env_get
+from mindspore.ops.function.clip_func import clip_by_global_norm
+from mindspore.ops.composite.multitype_ops.add_impl import hyper_add
+from mindspore.ops.composite.multitype_ops.ones_like_impl import ones_like
+from mindspore.ops.composite.multitype_ops.zeros_like_impl import zeros_like
+from mindspore.ops.function.random_func import normal, laplace, uniform, gamma, poisson, multinomial
+from mindspore.ops.composite.math_ops import matmul, cummin, mm
+from mindspore.ops.function.math_func import count_nonzero, tensor_dot, vecdot, dot, batch_dot
+from mindspore.ops.function.array_func import repeat_interleave, repeat_elements, sequence_mask
+from mindspore.ops.function.vmap_func import _VmapGeneralPreprocess, _VmapGeneralRule
+from mindspore.ops.function.clip_func import clip_by_value
+
+
+__all__ = [
+    'env_get',
+    'add_flags',
+    'tail',
+    'Map',
+    'MultitypeFuncGraph',
+    'GradOperation',
+    'HyperMap',
+    'hyper_add',
+    'zeros_like',
+    'ones_like',
+    'zip_operation',
+    'normal',
+    'laplace',
+    'uniform',
+    'gamma',
+    'poisson',
+    'multinomial',
+    'cummin',
+    'repeat_elements',
+    'repeat_interleave',
+    'sequence_mask',
+    'matmul',
+    'mm',
+    'count_nonzero',
+    'tensor_dot',
+    'vecdot',
+    'dot',
+    'batch_dot',
+    '_Grad',
+    '_Vmap',
+    'iter_converter',
+    'ms_hasnext',
+    'ms_next']