mindspore 2.2.10__cp39-cp39-win_amd64.whl → 2.2.14__cp39-cp39-win_amd64.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.

Potentially problematic release.


This version of mindspore might be problematic. Click here for more details.

Files changed (128) hide show
  1. mindspore/.commit_id +1 -1
  2. mindspore/__init__.py +2 -1
  3. mindspore/_c_dataengine.cp39-win_amd64.pyd +0 -0
  4. mindspore/_c_expression.cp39-win_amd64.pyd +0 -0
  5. mindspore/_c_mindrecord.cp39-win_amd64.pyd +0 -0
  6. mindspore/_extends/parse/__init__.py +3 -2
  7. mindspore/_extends/parse/parser.py +6 -1
  8. mindspore/_extends/parse/standard_method.py +12 -2
  9. mindspore/common/_utils.py +16 -0
  10. mindspore/common/tensor.py +0 -2
  11. mindspore/communication/management.py +3 -0
  12. mindspore/context.py +34 -4
  13. mindspore/dataset/engine/cache_client.py +8 -5
  14. mindspore/dataset/engine/datasets.py +23 -0
  15. mindspore/dataset/engine/validators.py +1 -1
  16. mindspore/dataset/vision/py_transforms_util.py +2 -2
  17. mindspore/dnnl.dll +0 -0
  18. mindspore/experimental/optim/lr_scheduler.py +5 -6
  19. mindspore/jpeg62.dll +0 -0
  20. mindspore/mindrecord/tools/cifar100_to_mr.py +49 -57
  21. mindspore/mindrecord/tools/cifar10_to_mr.py +46 -55
  22. mindspore/mindrecord/tools/csv_to_mr.py +3 -8
  23. mindspore/mindrecord/tools/mnist_to_mr.py +4 -9
  24. mindspore/mindrecord/tools/tfrecord_to_mr.py +1 -4
  25. mindspore/mindspore_backend.dll +0 -0
  26. mindspore/mindspore_common.dll +0 -0
  27. mindspore/mindspore_core.dll +0 -0
  28. mindspore/mindspore_glog.dll +0 -0
  29. mindspore/mindspore_shared_lib.dll +0 -0
  30. mindspore/nn/layer/activation.py +1 -1
  31. mindspore/nn/layer/embedding.py +2 -2
  32. mindspore/nn/layer/flash_attention.py +48 -135
  33. mindspore/nn/loss/loss.py +1 -1
  34. mindspore/nn/optim/ada_grad.py +2 -2
  35. mindspore/nn/optim/sgd.py +3 -2
  36. mindspore/nn/wrap/__init__.py +4 -2
  37. mindspore/nn/wrap/cell_wrapper.py +6 -3
  38. mindspore/numpy/math_ops.py +1 -1
  39. mindspore/opencv_core452.dll +0 -0
  40. mindspore/opencv_imgcodecs452.dll +0 -0
  41. mindspore/opencv_imgproc452.dll +0 -0
  42. mindspore/ops/__init__.py +3 -0
  43. mindspore/ops/_grad_experimental/grad_array_ops.py +0 -31
  44. mindspore/ops/_grad_experimental/grad_comm_ops.py +4 -2
  45. mindspore/ops/_grad_experimental/grad_inner_ops.py +8 -0
  46. mindspore/ops/_grad_experimental/grad_math_ops.py +37 -17
  47. mindspore/ops/_op_impl/aicpu/__init__.py +1 -0
  48. mindspore/ops/_op_impl/aicpu/generate_eod_mask.py +38 -0
  49. mindspore/ops/_op_impl/aicpu/linear_sum_assignment.py +21 -2
  50. mindspore/ops/function/array_func.py +6 -5
  51. mindspore/ops/function/debug_func.py +1 -1
  52. mindspore/ops/function/linalg_func.py +21 -11
  53. mindspore/ops/function/math_func.py +3 -0
  54. mindspore/ops/function/nn_func.py +13 -11
  55. mindspore/ops/function/parameter_func.py +2 -0
  56. mindspore/ops/function/sparse_unary_func.py +2 -2
  57. mindspore/ops/function/vmap_func.py +1 -0
  58. mindspore/ops/operations/__init__.py +5 -2
  59. mindspore/ops/operations/_embedding_cache_ops.py +1 -1
  60. mindspore/ops/operations/_grad_ops.py +3 -4
  61. mindspore/ops/operations/_inner_ops.py +56 -1
  62. mindspore/ops/operations/_quant_ops.py +4 -4
  63. mindspore/ops/operations/_rl_inner_ops.py +1 -1
  64. mindspore/ops/operations/array_ops.py +15 -4
  65. mindspore/ops/operations/custom_ops.py +1 -1
  66. mindspore/ops/operations/debug_ops.py +1 -1
  67. mindspore/ops/operations/image_ops.py +3 -3
  68. mindspore/ops/operations/inner_ops.py +49 -0
  69. mindspore/ops/operations/math_ops.py +65 -3
  70. mindspore/ops/operations/nn_ops.py +95 -28
  71. mindspore/ops/operations/random_ops.py +2 -0
  72. mindspore/ops/operations/sparse_ops.py +4 -4
  73. mindspore/ops/silent_check.py +162 -0
  74. mindspore/parallel/__init__.py +3 -2
  75. mindspore/parallel/_auto_parallel_context.py +82 -3
  76. mindspore/parallel/_parallel_serialization.py +34 -2
  77. mindspore/parallel/_tensor.py +3 -1
  78. mindspore/parallel/_transformer/transformer.py +8 -8
  79. mindspore/parallel/checkpoint_transform.py +191 -45
  80. mindspore/profiler/parser/ascend_cluster_generator.py +111 -0
  81. mindspore/profiler/parser/ascend_communicate_generator.py +315 -0
  82. mindspore/profiler/parser/ascend_flops_generator.py +8 -2
  83. mindspore/profiler/parser/ascend_fpbp_generator.py +8 -2
  84. mindspore/profiler/parser/ascend_hccl_generator.py +2 -2
  85. mindspore/profiler/parser/ascend_msprof_exporter.py +30 -6
  86. mindspore/profiler/parser/ascend_msprof_generator.py +16 -5
  87. mindspore/profiler/parser/ascend_op_generator.py +15 -7
  88. mindspore/profiler/parser/ascend_timeline_generator.py +5 -2
  89. mindspore/profiler/parser/base_timeline_generator.py +11 -3
  90. mindspore/profiler/parser/cpu_gpu_timeline_generator.py +2 -1
  91. mindspore/profiler/parser/framework_parser.py +8 -2
  92. mindspore/profiler/parser/memory_usage_parser.py +8 -2
  93. mindspore/profiler/parser/minddata_analyzer.py +8 -2
  94. mindspore/profiler/parser/minddata_parser.py +1 -1
  95. mindspore/profiler/parser/msadvisor_analyzer.py +4 -2
  96. mindspore/profiler/parser/msadvisor_parser.py +9 -3
  97. mindspore/profiler/profiling.py +97 -25
  98. mindspore/rewrite/api/node.py +1 -1
  99. mindspore/rewrite/api/symbol_tree.py +2 -2
  100. mindspore/rewrite/parsers/for_parser.py +6 -6
  101. mindspore/rewrite/parsers/module_parser.py +4 -4
  102. mindspore/tinyxml2.dll +0 -0
  103. mindspore/train/callback/_checkpoint.py +8 -8
  104. mindspore/train/callback/_landscape.py +2 -3
  105. mindspore/train/callback/_summary_collector.py +6 -7
  106. mindspore/train/dataset_helper.py +6 -0
  107. mindspore/train/model.py +17 -5
  108. mindspore/train/serialization.py +6 -1
  109. mindspore/train/summary/_writer_pool.py +1 -1
  110. mindspore/train/summary/summary_record.py +5 -6
  111. mindspore/turbojpeg.dll +0 -0
  112. mindspore/version.py +1 -1
  113. {mindspore-2.2.10.dist-info → mindspore-2.2.14.dist-info}/METADATA +3 -2
  114. {mindspore-2.2.10.dist-info → mindspore-2.2.14.dist-info}/RECORD +117 -124
  115. mindspore/ops/_op_impl/_custom_op/flash_attention/__init__.py +0 -0
  116. mindspore/ops/_op_impl/_custom_op/flash_attention/attention.py +0 -406
  117. mindspore/ops/_op_impl/_custom_op/flash_attention/constants.py +0 -41
  118. mindspore/ops/_op_impl/_custom_op/flash_attention/flash_attention_bwd.py +0 -467
  119. mindspore/ops/_op_impl/_custom_op/flash_attention/flash_attention_fwd.py +0 -563
  120. mindspore/ops/_op_impl/_custom_op/flash_attention/flash_attention_impl.py +0 -193
  121. mindspore/ops/_op_impl/_custom_op/flash_attention/tik_ops_utils.py +0 -435
  122. mindspore/ops/_op_impl/_custom_op/flash_attention/tiling_strategy/__init__.py +0 -0
  123. mindspore/ops/_op_impl/_custom_op/flash_attention/tiling_strategy/sparse_tiling.py +0 -45
  124. mindspore/ops/_op_impl/_custom_op/flash_attention/tiling_strategy/strategy.py +0 -67
  125. mindspore/ops/_op_impl/_custom_op/flash_attention/tiling_strategy/wukong_tiling.py +0 -62
  126. {mindspore-2.2.10.dist-info → mindspore-2.2.14.dist-info}/WHEEL +0 -0
  127. {mindspore-2.2.10.dist-info → mindspore-2.2.14.dist-info}/entry_points.txt +0 -0
  128. {mindspore-2.2.10.dist-info → mindspore-2.2.14.dist-info}/top_level.txt +0 -0
@@ -227,7 +227,11 @@ class ExpandDims(PrimitiveWithCheck):
227
227
  def infer_value(self, input_x, axis):
228
228
  value = None
229
229
  if input_x is not None and axis is not None:
230
- value = Tensor(np.expand_dims(input_x.asnumpy(), axis))
230
+ dtype = input_x.dtype
231
+ if input_x.dtype == mstype.bfloat16:
232
+ cpu_cast = Cast().set_device("CPU")
233
+ input_x = cpu_cast(input_x, mstype.float32)
234
+ value = Tensor(np.expand_dims(input_x.asnumpy(), axis), dtype)
231
235
  return value
232
236
 
233
237
 
@@ -375,6 +379,9 @@ class Cast(PrimitiveWithCheck):
375
379
  if isinstance(x, (int, float)):
376
380
  value = Tensor(np.array(x).astype(np_dst_type), dtype=dst_type)
377
381
  else:
382
+ if x.dtype == mstype.bfloat16:
383
+ cpu_cast = Cast().set_device("CPU")
384
+ x = cpu_cast(x, mstype.float32)
378
385
  value = Tensor(x.asnumpy().astype(np_dst_type), dtype=dst_type)
379
386
  return value
380
387
 
@@ -5837,6 +5844,9 @@ class SpaceToBatchND(Primitive):
5837
5844
  ``Ascend`` ``GPU`` ``CPU``
5838
5845
 
5839
5846
  Examples:
5847
+ >>> import mindspore
5848
+ >>> from mindspore import Tensor, ops
5849
+ >>> import numpy as np
5840
5850
  >>> block_shape = [2, 2]
5841
5851
  >>> paddings = [[0, 0], [0, 0]]
5842
5852
  >>> space_to_batch_nd = ops.SpaceToBatchND(block_shape, paddings)
@@ -6279,8 +6289,8 @@ class Sort(Primitive):
6279
6289
  Sorts the elements of the input tensor along the given dimension in the specified order.
6280
6290
 
6281
6291
  .. warning::
6282
- Currently, the data types of Float16 is well supported.
6283
- Using Float32 might cause loss of accuracy.
6292
+ Currently, the data types of float16, uint8, int8, int16, int32, int64 are well supported.
6293
+ If use float32, it may cause loss of accuracy.
6284
6294
 
6285
6295
  Args:
6286
6296
  axis (int, optional): The dimension to sort along. Default: ``-1``, means the last dimension.
@@ -8719,7 +8729,8 @@ class TopK(Primitive):
8719
8729
  - GPU: float16, float32.
8720
8730
  - CPU: all numeric types.
8721
8731
 
8722
- - **k** (int) - The number of top elements to be computed along the last dimension, constant input is needed.
8732
+ - **k** (Union(Tensor, int)) - The number of top elements to be computed along the last dimension.
8733
+ If `k` is a Tensor, the supported dtype is int32 and it should be 0-D or 1-D with shape :math:`(1, )` .
8723
8734
 
8724
8735
  Outputs:
8725
8736
  A tuple consisting of `values` and `indexes`.
@@ -470,7 +470,7 @@ class Custom(ops.PrimitiveWithInfer):
470
470
  op_path_in_cache = [] # Save paths for op functions created in the cached.
471
471
  custom_aot_warning = True # Flag to enable warnings about custom aot path white list
472
472
 
473
- def __init__(self, func, out_shape=None, out_dtype=None, func_type=HYBRID_TYPE, bprop=None, reg_info=None):
473
+ def __init__(self, func, out_shape=None, out_dtype=None, func_type="hybrid", bprop=None, reg_info=None):
474
474
  super().__init__("Custom")
475
475
 
476
476
  self.supported_targets = [ASCEND, GPU, CPU]
@@ -472,7 +472,7 @@ class Print(Primitive):
472
472
 
473
473
  Examples:
474
474
  >>> import numpy as np
475
- >>> from mindspore import Tensor, nn
475
+ >>> from mindspore import Tensor, nn, ops
476
476
  >>> class PrintDemo(nn.Cell):
477
477
  ... def __init__(self):
478
478
  ... super(PrintDemo, self).__init__()
@@ -388,7 +388,7 @@ class NonMaxSuppressionV3(Primitive):
388
388
  single score associated with each box (i.e., each row of the `boxes` Tensor).
389
389
  It is required that the number of scores in `scores` must be equal to the number of boxes in `boxes`.
390
390
  The supported data type is float32.
391
- - **max_output_size** (Union[Tensor, Number.Int]) - A scalar integer Tensor representing the maximum
391
+ - **max_output_size** (Union[Tensor, Number.int]) - A scalar integer Tensor representing the maximum
392
392
  number of boxes to be selected by non max suppression. The supported data type is int32.
393
393
  - **iou_threshold** (Union[Tensor, Number.Float]) - A scalar float Tensor represents the threshold
394
394
  used for determining if the intersection over union (IOU) between boxes is too high.
@@ -459,7 +459,7 @@ class NonMaxSuppressionWithOverlaps(Primitive):
459
459
  single score associated with each box (i.e., each row of the `boxes` Tensor).
460
460
  It is required that the number of scores in `scores` must be equal to the number of boxes in `boxes`.
461
461
  The supported data type is float32.
462
- - **max_output_size** (Union[Tensor, Number.Int]) - A scalar integer Tensor representing the maximum
462
+ - **max_output_size** (Union[Tensor, Number.int]) - A scalar integer Tensor representing the maximum
463
463
  number of boxes to be selected by non max suppression, and max_output_size must be equal to or greater
464
464
  than 0.
465
465
  Types allowed:int32.
@@ -816,7 +816,7 @@ class ResizeBicubic(Primitive):
816
816
  Examples:
817
817
  >>> import mindspore
818
818
  >>> import numpy as np
819
- >>> from mindspore import Tensor, ops
819
+ >>> from mindspore import Tensor, ops, nn
820
820
  >>> class NetResizeBicubic(nn.Cell):
821
821
  ... def __init__(self):
822
822
  ... super(NetResizeBicubic, self).__init__()
@@ -642,6 +642,55 @@ class FusedAdaFactorWithGlobalNorm(FusedAdaFactor):
642
642
  return param_type
643
643
 
644
644
 
645
+ class GenerateEodMask(Primitive):
646
+ r"""
647
+ Given the input `inputs_ids`, if found eod_token_id, the output position and attention mask matrix will be reset.
648
+ This means the `position_id` will start counting from 0, and the corresponding mask matrix will be filled with 0.
649
+
650
+ Args:
651
+ eod_token_id (int) - In the NLP scenario, this value corresponds to the id of
652
+ the symbol of 'EodOfDocument' in the vocabulary.
653
+
654
+ Inputs:
655
+ - **inputs_ids** (Tensor) - token id, a 2-D Tensor with shape :math:`(batch\_size, seq\_length)`.
656
+
657
+ Outputs:
658
+ - **position_id** (Tensor) - position id matrix with same shape and type as original `inputs_ids`.
659
+ - **attention_mask** (Tensor) - attention mask matrix with type
660
+ float16 and shape :math:`(batch\_size, seq\_length)`.
661
+
662
+ Supported Platforms:
663
+ ``Ascend``
664
+
665
+ Examples:
666
+ >>> op = ops.GenerateEodMask(eod_token_id=0)
667
+ >>> position, mask = op(Tensor([[1, 0, 3], [1, 0, 0]], dtype=mindspore.int32))
668
+ >>> print(position)
669
+ [[0 1 0] [0 0 1]]
670
+ >>> print(mask)
671
+ [[[ 1. 0. 0.]
672
+ [1. 1. 0.]
673
+ [0. 0. 1.]]
674
+ [[1. 0. 0.]
675
+ [0. 1. 0.]
676
+ [0. 1. 1.]]]
677
+
678
+ Raises:
679
+ - **TypeError** - If `eod_token_id` is not int.
680
+ - **TypeError** - If `inputs_ids` is not int.
681
+ - **ValueError** - If `inputs_ids` is not a 2-D Tensor.
682
+ """
683
+ @prim_attr_register
684
+ def __init__(self, n_pos, eod_token_id, n_step, n_error_mode='specific'):
685
+ """Initialize GenerateEodMask"""
686
+ validator.check_value_type("eod_token_id", eod_token_id, [int], self.name)
687
+ validator.check_value_type("n_pos", n_pos, [int], self.name)
688
+ validator.check_value_type("n_step", n_step, [list], self.name)
689
+ validator.check_value_type("n_error_mode", n_error_mode, [str], self.name)
690
+ self.init_prim_io_names(inputs=['inputs_ids'],
691
+ outputs=['position_ids'])
692
+
693
+
645
694
  class ScaleGrad(PrimitiveWithInfer):
646
695
  """
647
696
  Scale the input grad according to the loss scale.
@@ -123,6 +123,64 @@ class _MathBinaryOp(_BinaryOp):
123
123
  real_shape = [dim if cmp_dim > 0 else cmp_dim for dim, cmp_dim in zip(shape_value, cmp_shape)]
124
124
  return tuple(real_shape)
125
125
 
126
+ class SilentCheck(Primitive):
127
+ """
128
+ Implement SilentCheck on `pre_val`, `min_val`, `max_val`, `result` and
129
+ update them inplace with given parameters.
130
+
131
+ Args:
132
+ c_min_steps (int): an int determines...
133
+
134
+ c_thresh_l1 (float): a float determines...
135
+
136
+ c_coeff_l1 (float): a float determines...
137
+
138
+ c_thresh_l2 (float): a float determines...
139
+
140
+ c_coeff_l2 (float): a float determines...
141
+
142
+ Inputs:
143
+ - **val** (Tensor) - Tensor with dtype float32.
144
+ - **input_grad** (Parameter) - Tensor with dtype float32.
145
+ - **pre_val** (Parameter) - Input Parameter with dtype float32.
146
+ - **min_val** (Parameter) - Input Parameter with dtype float32.
147
+ - **max_val** (Parameter) - Input Parameter with dtype float32.
148
+ - **val_counter** (Parameter) - Input Parameter with dtype int32.
149
+
150
+ Outputs:
151
+ Tuple of 5 Tensors, the updated parameters.
152
+ - **input_grad** (Tensor) - Tensor with dtype float32.
153
+ - **pre_val** (Tensor) - Tensor with dtype float32.
154
+ - **min_val** (Tensor) - Tensor with dtype float32.
155
+ - **max_val** (Tensor) - Tensor with dtype float32.
156
+ - **result** (Tensor) - Tensor with dtype int32.
157
+
158
+ Raises:
159
+ TypeError: If `val` is not Tensor with dtype float32.
160
+ TypeError: If `result` is not Tensor with dtype int32.
161
+ TypeError: If `pre_val`, `min_val`, `max_val`, `input_grad` are not all Parameter type with dtype float32.
162
+ TypeError: If `c_thresh_l1` or `c_coeff_l1` is not a float number.
163
+ TypeError: If `c_min_steps` is not an int number.
164
+
165
+ Supported Platforms:
166
+ ``Ascend``
167
+
168
+ Examples:
169
+ >>> from mindspore.ops.operations.math_ops import SilentCheck
170
+ >>> silent_check = SilentCheck()
171
+ xxx
172
+ """
173
+
174
+ @prim_attr_register
175
+ def __init__(self, c_min_steps, c_thresh_l1, c_coeff_l1, c_thresh_l2, c_coeff_l2):
176
+ """Initialize SilentCheck."""
177
+ validator.check_value_type("c_min_steps", c_min_steps, [int], self.name)
178
+ validator.check_value_type("c_thresh_l1", c_thresh_l1, [float], self.name)
179
+ validator.check_value_type("c_coeff_l1", c_coeff_l1, [float], self.name)
180
+ validator.check_value_type("c_thresh_l2", c_thresh_l2, [float], self.name)
181
+ validator.check_value_type("c_coeff_l2", c_coeff_l2, [float], self.name)
182
+ self.add_prim_attr('side_effect_mem', True)
183
+
126
184
 
127
185
  class _BitwiseBinaryOp(_MathBinaryOp):
128
186
  """
@@ -462,6 +520,7 @@ class AssignAdd(Primitive):
462
520
  >>> import mindspore
463
521
  >>> import numpy as np
464
522
  >>> from mindspore import Tensor, ops, nn
523
+ >>> from mindspore.common.initializer import initializer
465
524
  >>> class Net(nn.Cell):
466
525
  ... def __init__(self):
467
526
  ... super(Net, self).__init__()
@@ -512,6 +571,7 @@ class AssignSub(Primitive):
512
571
  >>> import mindspore
513
572
  >>> import numpy as np
514
573
  >>> from mindspore import Tensor, ops, nn
574
+ >>> from mindspore.common.initializer import initializer
515
575
  >>> class Net(nn.Cell):
516
576
  ... def __init__(self):
517
577
  ... super(Net, self).__init__()
@@ -6569,9 +6629,9 @@ class LinSpace(Primitive):
6569
6629
 
6570
6630
  Inputs:
6571
6631
  - **start** (Tensor) - Start value of interval, 0-D Tensor with dtype float32 or float64.
6572
- - **stop** (Tensor) - Last value of interval, 0-D Tensor with dtype float32 or float64.
6573
- - **num** (int) - Number of ticks in the interval, inclusive of `start` and `stop`.
6574
- Supported dtypes: int32, int64.
6632
+ - **stop** (Tensor) - Last value of interval, 0-D Tensor with dtype float32 or float64.
6633
+ - **num** (Union[int, Tensor]) - Number of ticks in the interval, inclusive of `start` and `stop`.
6634
+ Must be a positive integer. When the input is Tensor, it must be a 0-D Tensor with dtype int32 or int64.
6575
6635
 
6576
6636
  Outputs:
6577
6637
  Tensor, has the same shape and dtype as `start`.
@@ -7253,6 +7313,7 @@ class Igamma(Primitive):
7253
7313
 
7254
7314
  Examples:
7255
7315
  >>> import numpy as np
7316
+ >>> import mindspore
7256
7317
  >>> from mindspore import Tensor, ops
7257
7318
  >>> a = Tensor(np.array([2.0, 4.0, 6.0, 8.0]).astype(np.float32))
7258
7319
  >>> x = Tensor(np.array([2.0, 3.0, 4.0, 5.0]).astype(np.float32))
@@ -7291,6 +7352,7 @@ class Igammac(Primitive):
7291
7352
  ``Ascend`` ``GPU`` ``CPU``
7292
7353
 
7293
7354
  Examples:
7355
+ >>> import mindspore
7294
7356
  >>> import numpy as np
7295
7357
  >>> from mindspore import Tensor, ops
7296
7358
  >>> a = Tensor(np.array([2.0, 4.0, 6.0, 8.0]).astype(np.float32))
@@ -3777,7 +3777,7 @@ class LayerNorm(Primitive):
3777
3777
  - **output_x** (Tensor) - The normalized input, has the same type and shape as the `input_x`.
3778
3778
  - **mean** (Tensor) - The first `begin_norm_axis` dimensions of `mean` shape is the same as `input_x`,
3779
3779
  and the remaining dimensions are 1. Suppose the shape of the `input_x` is :math:`(x_1, x_2, \ldots, x_R)`,
3780
- the shape of the `mean` is :math:`(x_1, \ldots, x_{begin_params_axis}, 1, \ldots, 1)`
3780
+ the shape of the `mean` is :math:`(x_1, \ldots, x_{begin\_params\_axis}, 1, \ldots, 1)`
3781
3781
  (when `begin_params_axis=0`, the shape of `mean` is :math:`(1, \ldots, 1)` ).
3782
3782
  - **variance** (Tensor) - Shape is the same as `mean` .
3783
3783
 
@@ -4917,6 +4917,7 @@ class Adam(Primitive):
4917
4917
  >>> import mindspore
4918
4918
  >>> import numpy as np
4919
4919
  >>> from mindspore import Tensor, nn, ops
4920
+ >>> from mindspore import Parameter
4920
4921
  >>> class Net(nn.Cell):
4921
4922
  ... def __init__(self):
4922
4923
  ... super(Net, self).__init__()
@@ -9991,6 +9992,9 @@ class FractionalMaxPool3DWithFixedKsize(Primitive):
9991
9992
  ``Ascend`` ``GPU`` ``CPU``
9992
9993
 
9993
9994
  Examples:
9995
+ >>> import numpy as np
9996
+ >>> from mindspore import Tensor, ops
9997
+ >>> from mindspore import dtype as mstype
9994
9998
  >>> x = Tensor(np.array([1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16])
9995
9999
  ... .reshape([1, 1, 2, 2, 4]), mstype.float32)
9996
10000
  >>> random_samples = Tensor(np.array([0.7, 0.7, 0.7]).reshape([1, 1, 3]), mstype.float32)
@@ -11363,7 +11367,7 @@ class PromptFlashAttention(Primitive):
11363
11367
  For each element, 0 indicates retention and 1 indicates discard. Input tensor of shape :math:`(B, 1, S, S)`.
11364
11368
  - **actual_seq_lengths** (Tensor): Describe actual sequence length of each input with data type of int.
11365
11369
  - **actual_seq_lengths_kv** (Tensor): Describe actual sequence length of each input with data type of int.
11366
- - **padding_mask** (Tensor) - The padding mask tensor with data type of float16 or float32
11370
+ - **pse_shift** (Tensor) - The position encoding tensor with data type of float16 or float32.
11367
11371
  - **dep_scale1** (Tensor)
11368
11372
  - **quant_scale1** (Tensor)
11369
11373
  - **deq_scale2** (Tensor)
@@ -11406,7 +11410,7 @@ class PromptFlashAttention(Primitive):
11406
11410
  validator.check_value_type('num_key_value_heads', num_key_value_heads, [int], self.name)
11407
11411
  validator.check_value_type('sparse_mode', sparse_mode, [int], self.name)
11408
11412
  self.init_prim_io_names(inputs=["query", "key", "value", "attn_mask", "actual_seq_lengths",
11409
- "actual_seq_lengths_kv", "padding_mask", "deq_scale1", "quant_scale1",
11413
+ "actual_seq_lengths_kv", "pse_shift", "deq_scale1", "quant_scale1",
11410
11414
  "deq_scale2", "quant_scale2", "quant_offset2"],
11411
11415
  outputs=["attention_out"])
11412
11416
 
@@ -11417,41 +11421,50 @@ class FlashAttentionScore(Primitive):
11417
11421
  .. warning::
11418
11422
  This is an experimental API that is subject to change or deletion.
11419
11423
  B -- Batch size
11420
- S -- Sequence length
11421
- H -- Hidden size
11422
- N -- Num heads
11423
- D -- Dim size
11424
+ S1 -- Sequence length of query
11425
+ S2 -- Sequence length of key and value
11426
+ N1 -- Num heads of query
11427
+ N2 -- Num heads of key and value, and N2 must be a factor of N1
11428
+ D -- head size
11429
+ H1 -- Hidden size of query, which equals to N1 * D
11430
+ H2 -- Hidden size of key and value, which equals to N2 * D
11424
11431
  Args:
11425
- head_num (int): The number of the heads.
11432
+ head_num (int): The head num of query.
11426
11433
  keep_prob (float): The keep probability of dropout. Default: 1.0.
11427
11434
  scale_value (float): The scale value. Default: 1.0.
11428
11435
  pre_tokens (int): Previous tokens. Default: 65536.
11429
11436
  next_tokens (int): Next tokens. Default: 65536.
11430
11437
  inner_precise (int): Specify the execution mode, where 0 indicates high precision mode and 1 indicates high
11431
- performance mode. Default: 0.
11438
+ performance mode. Only support 0 currently. Default: 0.
11432
11439
  input_layout (str, optional): Specifies the layout of `query`, the value must be one of ["BSH", "BNSD"].
11433
11440
  Default: "BSH".
11434
11441
  sparse_mode (int): Default 0.
11435
11442
 
11436
11443
  Inputs:
11437
- - **query** (Tensor) - The query tensor with data type must be in [float16, float32, bfloat16].
11438
- Input tensor of shape :math:`(B, S, H)`.
11439
- - **key** (Tensor) - The key tensor with data must be in [float16, float32, bfloat16].
11440
- Input tensor of shape :math:`(B, S, H)`.
11441
- - **value** (Tensor) - The value tensor with data must be in [float16, float32, bfloat16].
11442
- Input tensor of shape :math:`(B, S, H)`.
11443
- - **attn_mask** (Tensor) - The attention mask tensor with data type of uint8.
11444
- For each element, 0 indicates retention and 1 indicates discard. Input tensor of shape :math:`(B, 1, S, S)`.
11445
- - **drop_mask** (Tensor) - The dropout mask tensor with data type of UInt8.
11446
- Input tensor of shape :math:`(B, N, S, S // 8) or ()`.
11447
- - **real_shift** (None) - The position embedding code of float16 or float32, not implemented yet.
11444
+ - **query** (Tensor[float16, float32, bfloat16]) - The query tensor.
11445
+ Input tensor of shape :math:`(B, S1, H1)` or `(B, N1, S1, D)`.
11446
+ - **key** (Tensor[float16, float32, bfloat16]) - The key tensor.
11447
+ Input tensor of shape :math:`(B, S2, H2)` or `(B, N2, S2, D)`.
11448
+ - **value** (Tensor[float16, float32, bfloat16]) - The value tensor.
11449
+ Input tensor of shape :math:`(B, S2, H2)` or `(B, N2, S2, D)`.
11450
+ - **real_shift** (Tensor[float16, float32, bfloat16], None) - The position embedding code.
11451
+ Input tensor of shape :math: `(B, N1, S1, S2)` or `(B, N1, 1, S2)`.
11452
+ - **drop_mask** (Tensor[uint8], None) - The dropout mask tensor.
11453
+ Input tensor of shape :math:`(B, N1, S1, S2 // 8) or None`.
11448
11454
  - **padding_mask** (None) - The padding mask of float16 or float32, not implemented yet.
11449
- - **prefix** (None) - Not implemented yet.
11455
+ - **attn_mask** (Tensor[uint8], None) - The attention mask tensor.
11456
+ For each element, 0 indicates retention and 1 indicates discard.
11457
+ Input tensor of shape :math:`(B, N1, S1, S2)`, `(B, 1, S1, S2)` or `(S1, S2)`.
11458
+ - **prefix** (Tensor[int64], None) - Not implemented yet.
11459
+ Input tensor of shape :math:`(B,)`.
11450
11460
 
11451
11461
  Outputs:
11452
- - **attention_out** (Tensor) - (B, S, H)
11453
- - **softmax_max** (Tensor) - (B, N, S, 16)/(B, N, S, 8) when fp16/fp32
11454
- - **softmax_sum** (Tensor) - (B, N, S, 16)/(B, N, S, 8) when fp16/fp32
11462
+ - **softmax_max** (Tensor[float32]) - (B, N1, S1, 8)
11463
+ - **softmax_sum** (Tensor[float32]) - (B, N1, S1, 8)
11464
+ - **softmax_out** (Tensor[float32]) - Useless output, ignore it. Output tensor of shape : `()`
11465
+ - **attention_out** (Tensor[float16, float32, bfloat16]) - The output of attention, its shape, and data type
11466
+ are the same as the query.
11467
+
11455
11468
  Supported Platforms:
11456
11469
  ``Ascend``
11457
11470
  """
@@ -11469,14 +11482,14 @@ class FlashAttentionScore(Primitive):
11469
11482
  validator.check_value_type('next_tokens', next_tokens, [int], self.name)
11470
11483
  validator.check_value_type('inner_precise', inner_precise, [int], self.name)
11471
11484
  validator.check_value_type('sparse_mode', sparse_mode, [int], self.name)
11472
- if inner_precise not in [0, 1]:
11473
- raise ValueError(f"Attribute 'inner_precise' must be either 0 or 1, but got {inner_precise}")
11485
+ if inner_precise not in [0]:
11486
+ raise ValueError(f"Attribute 'inner_precise' must be 0, but got {inner_precise}")
11474
11487
  validator.check_value_type('input_layout', input_layout, [str], self.name)
11475
11488
  if input_layout not in ["BSH", "BNSD"]:
11476
11489
  raise ValueError(f"Attribute 'input_layout' must be either 'BSH' or 'BNSD', but got {input_layout}")
11477
11490
  self.init_prim_io_names(
11478
- inputs=['query', 'key', 'value', 'attn_mask', 'drop_mask', 'real_shift', 'padding_mask', 'prefix'],
11479
- outputs=['attention_out', 'softmax_max', 'softmax_sum'])
11491
+ inputs=['query', 'key', 'value', 'real_shift', 'drop_mask', 'padding_mask', 'attn_mask', 'prefix'],
11492
+ outputs=['softmax_max', 'softmax_sum', 'softmax_out', 'attention_out'])
11480
11493
 
11481
11494
 
11482
11495
  class RmsNorm(Primitive):
@@ -11514,3 +11527,57 @@ class RmsNorm(Primitive):
11514
11527
  """Initialize Dense."""
11515
11528
  validator.check_value_type("epsilon", epsilon, [float], self.name)
11516
11529
  self.init_prim_io_names(inputs=['x', 'gamma'], outputs=["y", "rstd"])
11530
+
11531
+
11532
+ class PagedAttention(Primitive):
11533
+ r"""
11534
+ .. warning::
11535
+ This is an experimental API that is subject to change or deletion.
11536
+ """
11537
+ @prim_attr_register
11538
+ def __init__(self, head_num, scale_value=1.0, kv_head_num=0):
11539
+ """Initialize PagedAttention"""
11540
+ validator.check_value_type('head_num', head_num, [int], self.name)
11541
+ validator.check_value_type('scale_value', scale_value, [float], self.name) # scale after qkbmm
11542
+ validator.check_value_type('kv_head_num', kv_head_num, [int], self.name) # for MQA
11543
+ self.init_prim_io_names(
11544
+ inputs=['query', 'key_cache', 'value_cache', 'block_tables', 'context_lens'],
11545
+ outputs=['attention_out'])
11546
+
11547
+
11548
+ class PagedAttentionMask(Primitive):
11549
+ r"""
11550
+ .. warning::
11551
+ This is an experimental API that is subject to change or deletion.
11552
+ """
11553
+ @prim_attr_register
11554
+ def __init__(self, head_num, scale_value=1.0, kv_head_num=0):
11555
+ """Initialize PagedAttentionMask"""
11556
+ validator.check_value_type('head_num', head_num, [int], self.name)
11557
+ validator.check_value_type('scale_value', scale_value, [float], self.name) # scale after qkbmm
11558
+ validator.check_value_type('kv_head_num', kv_head_num, [int], self.name) # for MQA
11559
+ self.init_prim_io_names(
11560
+ inputs=['query', 'key_cache', 'value_cache', 'block_tables', 'context_lens', 'alibi_mask'],
11561
+ outputs=['attention_out'])
11562
+
11563
+
11564
+ class ReshapeAndCache(Primitive):
11565
+ r"""
11566
+ .. warning::
11567
+ This is an experimental API that is subject to change or deletion.
11568
+ """
11569
+ __mindspore_signature__ = (
11570
+ sig.make_sig('key', dtype=sig.sig_dtype.T),
11571
+ sig.make_sig('value', dtype=sig.sig_dtype.T),
11572
+ sig.make_sig('key_cache', sig.sig_rw.RW_WRITE, dtype=sig.sig_dtype.T),
11573
+ sig.make_sig('value_cache', sig.sig_rw.RW_WRITE, dtype=sig.sig_dtype.T),
11574
+ sig.make_sig('slot_mapping', dtype=sig.sig_dtype.T1),
11575
+ )
11576
+
11577
+ @prim_attr_register
11578
+ def __init__(self):
11579
+ """Initialize ReshapeAndCache"""
11580
+ self.init_prim_io_names(
11581
+ inputs=['key', 'value', 'key_cache', 'value_cache', 'slot_mapping'],
11582
+ outputs=['key_out'])
11583
+ self.add_prim_attr('side_effect_mem', True)
@@ -1271,6 +1271,8 @@ class RandpermV2(Primitive):
1271
1271
  ``Ascend`` ``CPU``
1272
1272
 
1273
1273
  Examples:
1274
+ >>> from mindspore import Tensor, ops
1275
+ >>> from mindspore import dtype as mstype
1274
1276
  >>> n = Tensor([4], mstype.int64)
1275
1277
  >>> seed = 0
1276
1278
  >>> offset = 0
@@ -479,8 +479,8 @@ class SparseToDenseV2(Primitive):
479
479
  Tensor, converted from sparse tensor. The dtype is same as `values`, and the shape is `output_shape`.
480
480
 
481
481
  Raises:
482
- TypeError: If the dtype of `indices` is neither Int32 nor Int64.
483
- TypeError: If the dtype of `outputshape` is neither Int32 nor Int64.
482
+ TypeError: If the dtype of `indices` is neither int32 nor int64.
483
+ TypeError: If the dtype of `outputshape` is neither int32 nor int64.
484
484
  ValueError: If the shape of `output_shape`, shape of `indices`,
485
485
  shape of `default_value` and shape of `values` don't meet the parameter description.
486
486
  ValueError: If each Element of `output_shape` is not > 0.
@@ -2382,8 +2382,8 @@ class SparseCountSparseOutput(Primitive):
2382
2382
  Args:
2383
2383
  binary_output (bool) - If ``False`` , output the number of occurrences of each value,
2384
2384
  if ``True`` output 1 for orresponding values. Default: ``False`` .
2385
- minlength(Scalar) - Int type minimum value to count, Default: ``-1`` .
2386
- maxlength(Scalar) - Int type maximum value to count, Default: ``-1`` .
2385
+ minlength(Scalar) - int type minimum value to count, Default: ``-1`` .
2386
+ maxlength(Scalar) - int type maximum value to count, Default: ``-1`` .
2387
2387
 
2388
2388
  Inputs:
2389
2389
  - **indices** (Tensor) - Tensor representing the position of the element in the sparse
@@ -0,0 +1,162 @@
1
+ # Copyright 2024 Huawei Technologies Co., Ltd
2
+ #
3
+ # Licensed under the Apache License, Version 2.0 (the "License");
4
+ # you may not use this file except in compliance with the License.
5
+ # You may obtain a copy of the License at
6
+ #
7
+ # http://www.apache.org/licenses/LICENSE-2.0
8
+ #
9
+ # Unless required by applicable law or agreed to in writing, software
10
+ # distributed under the License is distributed on an "AS IS" BASIS,
11
+ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12
+ # See the License for the specific language governing permissions and
13
+ # limitations under the License.
14
+ # ============================================================================
15
+ """Silent Check."""
16
+ import os
17
+
18
+ from mindspore.common.tensor import Tensor
19
+ from mindspore.common.parameter import Parameter
20
+ import mindspore.common.dtype as mstype
21
+
22
+ from . import operations
23
+ from .operations._inner_ops import _MirrorSilentCheck
24
+ from .operations import RmsNorm as OriginRmsNorm
25
+ from .operations import LayerNorm as OriginLayerNorm
26
+ from .primitive import Primitive
27
+
28
+
29
+ NPU_ASD_ENABLE = 'NPU_ASD_ENABLE'
30
+
31
+
32
+ class ASDBase:
33
+ """
34
+ ASDBase is the base class of operator with accuracy-sensitive detection feature in python.
35
+
36
+ Args:
37
+ cls (Primitive): Original operator requiring accuracy-sensitive detection feature.
38
+ args (tuple): A variable parameter tuple to the original operator.
39
+ kwargs (dict): A variable parameter dictionary passed the original operator.
40
+
41
+ Supported Platforms:
42
+ ``Ascend``
43
+
44
+ Examples:
45
+ >>> from mindspore.ops.silent_check import ASDBase
46
+ >>> from mindspore.ops import LayerNorm as OriginLayerNorm
47
+ >>> class LayerNormASD(ASDBase):
48
+ ... def __init__(self, *args, **kwargs):
49
+ ... super().__init__(OriginLayerNorm, *args, **kwargs)
50
+ ... # init parameters for accuracy-sensitive detection by calling the base class method generate_params()
51
+ ... self.pre_val, self.min_val, self.max_val, self.cnt = self.generate_params()
52
+ ...
53
+ ... def __call__(self, input_x, gamma, beta):
54
+ ... if self.enable_check:
55
+ ... # execute accuracy-sensitive detection by calling the check_op of base class
56
+ ... input_x = self.check_op(
57
+ ... input_x, self.pre_val, self.min_val, self.max_val, self.cnt, None)
58
+ ... self.cnt += 1
59
+ ... # return the result of original operator
60
+ ... return self.op(input_x, gamma, beta)
61
+ """
62
+ _index = 0
63
+ __ms_class__ = True
64
+
65
+ def __init__(self, cls, *args, **kwargs):
66
+ self.op = cls(*args, **kwargs)
67
+ self.check_op = _MirrorSilentCheck()
68
+ self._suffix = "ASD_" + cls.__name__
69
+ primitive_attr = dir(Primitive)
70
+ self._op_attr_dict = {
71
+ name for name in primitive_attr if not name.startswith("_")}
72
+ self.enable_check = os.environ.get(NPU_ASD_ENABLE) == "1"
73
+
74
+ def __getattr__(self, name):
75
+ def method_wrapper(*args, **kwargs):
76
+ out = getattr(self.op, name)(*args, **kwargs)
77
+ if out is self.op:
78
+ return self
79
+ return out
80
+
81
+ if name in self._op_attr_dict:
82
+ if callable(getattr(self.op, name)):
83
+ return method_wrapper
84
+ if hasattr(self.op, name):
85
+ return getattr(self.op, name)
86
+ return super().__getattr__(self, name)
87
+
88
+ def __repr__(self):
89
+ return self.op.__repr__()
90
+
91
+ def generate_params(self):
92
+ """
93
+ Generate support params for accuracy-sensitive detection.
94
+
95
+ Returns:
96
+ tuple consisting of four elements.
97
+ The derived class initializes the parameters required for accuracy-sensitive detection by calling
98
+ this function.
99
+
100
+ Examples:
101
+ >>> from mindspore.ops.silent_check import ASDBase
102
+ >>> from mindspore.ops import LayerNorm as OriginLayerNorm
103
+ >>> class LayerNormASD(ASDBase):
104
+ ... def __init__(self, *args, **kwargs):
105
+ ... super().__init__(OriginLayerNorm, *args, **kwargs)
106
+ ... # init parameters for accuracy-sensitive detection by calling the base class function
107
+ ... self.pre_val, self.min_val, self.max_val, self.cnt = self.generate_params()
108
+ """
109
+ pre_val = Parameter(Tensor(0, mstype.float32),
110
+ name=f"{self._suffix}_pre_val_{self._index}",
111
+ requires_grad=False)
112
+ min_val = Parameter(Tensor(0, mstype.float32),
113
+ name=f"{self._suffix}_min_val_{self._index}",
114
+ requires_grad=False)
115
+ max_val = Parameter(Tensor(0, mstype.float32),
116
+ name=f"{self._suffix}_max_val_{self._index}",
117
+ requires_grad=False)
118
+ cnt = Parameter(Tensor(0, mstype.int32),
119
+ name=f"{self._suffix}_cnt_{self._index}",
120
+ requires_grad=False)
121
+ ASDBase._index += 1
122
+ return pre_val, min_val, max_val, cnt
123
+
124
+
125
+ class RmsNormASD(ASDBase):
126
+ """
127
+ RmsNorm with ASD.
128
+ """
129
+
130
+ def __init__(self, *args, **kwargs):
131
+ super().__init__(OriginRmsNorm, *args, **kwargs)
132
+ self.pre_val, self.min_val, self.max_val, self.cnt = self.generate_params()
133
+
134
+ def __call__(self, input_x, gamma):
135
+ if self.enable_check:
136
+ input_x = self.check_op(
137
+ input_x, self.pre_val, self.min_val, self.max_val, self.cnt, None)
138
+ self.cnt += 1
139
+ return self.op(input_x, gamma)
140
+
141
+
142
+ class LayerNormASD(ASDBase):
143
+ """
144
+ LayerNorm with ASD.
145
+ """
146
+
147
+ def __init__(self, *args, **kwargs):
148
+ super().__init__(OriginLayerNorm, *args, **kwargs)
149
+ self.pre_val, self.min_val, self.max_val, self.cnt = self.generate_params()
150
+
151
+ def __call__(self, input_x, gamma, beta):
152
+ if self.enable_check:
153
+ input_x = self.check_op(
154
+ input_x, self.pre_val, self.min_val, self.max_val, self.cnt, None)
155
+ self.cnt += 1
156
+ return self.op(input_x, gamma, beta)
157
+
158
+
159
+ def _silent_check():
160
+ if os.environ.get(NPU_ASD_ENABLE) == "1":
161
+ operations.LayerNorm = LayerNormASD
162
+ operations.RmsNorm = RmsNormASD