mindspore 2.4.0__cp39-cp39-manylinux1_x86_64.whl → 2.4.10__cp39-cp39-manylinux1_x86_64.whl

mindspore/.commit_id

@@ -1 +1 @@
-__commit_id__ = '[sha1]:8c86f33f,[branch]:(HEAD,origin/master,origin/HEAD,master)'
+__commit_id__ = '[sha1]:8e2ae935,[branch]:(HEAD,origin/r2.4.1,r2.4.1)'

mindspore/common/api.py

@@ -1703,7 +1703,6 @@ class _CellGraphExecutor:
         self._graph_executor = GraphExecutor_.get_instance()
         self._graph_executor.set_py_exe_path(sys.executable)
         self._graph_executor.set_kernel_build_server_dir(os.path.split(kernel_build_server.__file__)[0] + os.sep)
-        self._pid = os.getpid()
 
     def init_dataset(self, queue_name, dataset_size, batch_size, dataset_types, dataset_shapes,
                      input_indexs, phase='dataset', need_run=True):
@@ -1934,9 +1933,7 @@ class _CellGraphExecutor:
 
     def del_net_res(self, obj, net_id):
         """Clear the memory resource of a network."""
-        # no need to del net res by gc in independent dataset process which is a subprocess forked by main process
-        if self._pid == os.getpid():
-            self._graph_executor.del_net_res(obj, net_id)
+        self._graph_executor.del_net_res(obj, net_id)
 
     def _get_branch_control_input(self):
         if ('obf_ratio' not in self.obfuscate_config.keys()) or (

mindspore/common/file_system.py

@@ -22,6 +22,7 @@ class FileSystem:
         self.create_args = ("ab",)
         self.open = open
         self.open_args = ("rb",)
+        self.backend = "basic"
 
 
 def _register_basic_file_system(fs: FileSystem):
@@ -45,4 +46,5 @@ def _register_mindio_file_system(fs: FileSystem):
     fs.create_args = ()
     fs.open = mindio.open_file
     fs.open_args = ()
+    fs.backend = "mindio"
     return True

mindspore/common/initializer.py

@@ -103,6 +103,12 @@ def _numpy_seed():
     return np.random.randint(low=1, high=(1 << 63), dtype=np.int64)
 
 
+def _init_random_normal_inplace(mean, sigma, arr):
+    if sigma < 0:
+        raise ValueError("sigma < 0")
+    _random_normal(_numpy_seed(), arr, mean, sigma)
+
+
 def _init_random_normal(mean, sigma, shape):
     if sigma < 0:
         raise ValueError("sigma < 0")
@@ -111,12 +117,22 @@ def _init_random_normal(mean, sigma, shape):
     return data
 
 
+def _init_random_uniform_inplace(a, b, arr):
+    _random_uniform(_numpy_seed(), arr, a, b)
+
+
 def _init_random_uniform(a, b, shape):
     data = np.ndarray(shape=shape, dtype=np.float32)
     _random_uniform(_numpy_seed(), data, a, b)
     return data
 
 
+def _init_truncated_normal_inplace(a, b, mean, sigma, arr):
+    if sigma < 0:
+        raise ValueError("sigma < 0")
+    _truncated_normal(_numpy_seed(), arr, a, b, mean, sigma)
+
+
 def _init_truncated_normal(a, b, mean, sigma, shape):
     if sigma < 0:
         raise ValueError("sigma < 0")
@@ -298,9 +314,11 @@ class XavierNormal(Initializer):
         fan_in, fan_out = _calculate_fan_in_and_fan_out(arr.shape)
 
         std = self.gain * math.sqrt(2.0 / float(fan_in + fan_out))
-        data = _init_random_normal(0, std, arr.shape)
-
-        _assignment(arr, data)
+        if isinstance(arr, np.ndarray) and arr.dtype == np.float32:
+            _init_random_normal_inplace(0, std, arr)
+        else:
+            data = _init_random_normal(0, std, arr.shape)
+            _assignment(arr, data)
 
 
 @_register('xavier_uniform')
@@ -337,8 +355,11 @@ class XavierUniform(Initializer):
     def _initialize(self, arr):
         n_in, n_out = _calculate_fan_in_and_fan_out(arr.shape)
         boundary = self.gain * math.sqrt(6.0 / (n_in + n_out))
-        data = _init_random_uniform(-boundary, boundary, arr.shape)
-        _assignment(arr, data)
+        if isinstance(arr, np.ndarray) and arr.dtype == np.float32:
+            _init_random_uniform_inplace(-boundary, boundary, arr)
+        else:
+            data = _init_random_uniform(-boundary, boundary, arr.shape)
+            _assignment(arr, data)
 
 
 @_register('he_uniform')
@@ -386,8 +407,11 @@ class HeUniform(Initializer):
         gain = _calculate_gain(self.nonlinearity, self.negative_slope)
         std = gain / math.sqrt(fan)
         boundary = math.sqrt(3.0) * std
-        data = _init_random_uniform(-boundary, boundary, arr.shape)
-        _assignment(arr, data)
+        if isinstance(arr, np.ndarray) and arr.dtype == np.float32:
+            _init_random_uniform_inplace(-boundary, boundary, arr)
+        else:
+            data = _init_random_uniform(-boundary, boundary, arr.shape)
+            _assignment(arr, data)
 
 
 @_register('he_normal')
@@ -432,8 +456,11 @@ class HeNormal(Initializer):
         fan = _calculate_correct_fan(arr.shape, self.mode)
         gain = _calculate_gain(self.nonlinearity, self.negative_slope)
         std = gain / math.sqrt(fan)
-        data = _init_random_normal(0, std, arr.shape)
-        _assignment(arr, data)
+        if isinstance(arr, np.ndarray) and arr.dtype == np.float32:
+            _init_random_normal_inplace(0, std, arr)
+        else:
+            data = _init_random_normal(0, std, arr.shape)
+            _assignment(arr, data)
 
 
 class Constant(Initializer):
@@ -718,8 +745,11 @@ class Uniform(Initializer):
         self.scale = scale
 
     def _initialize(self, arr):
-        tmp = _init_random_uniform(-self.scale, self.scale, arr.shape)
-        _assignment(arr, tmp)
+        if isinstance(arr, np.ndarray) and arr.dtype == np.float32:
+            _init_random_uniform_inplace(-self.scale, self.scale, arr)
+        else:
+            tmp = _init_random_uniform(-self.scale, self.scale, arr.shape)
+            _assignment(arr, tmp)
 
 
 @_register()
@@ -749,8 +779,11 @@ class Normal(Initializer):
         self.mean = mean
 
     def _initialize(self, arr):
-        data = _init_random_normal(self.mean, self.sigma, arr.shape)
-        _assignment(arr, data)
+        if isinstance(arr, np.ndarray) and arr.dtype == np.float32:
+            _init_random_normal_inplace(self.mean, self.sigma, arr)
+        else:
+            data = _init_random_normal(self.mean, self.sigma, arr.shape)
+            _assignment(arr, data)
 
 
 @_register()
@@ -780,8 +813,11 @@ class TruncatedNormal(Initializer):
         self.b = b
 
     def _initialize(self, arr):
-        tmp = _init_truncated_normal(self.a, self.b, self.mean, self.sigma, arr.shape)
-        _assignment(arr, tmp)
+        if isinstance(arr, np.ndarray) and arr.dtype == np.float32:
+            _init_truncated_normal_inplace(self.a, self.b, self.mean, self.sigma, arr)
+        else:
+            tmp = _init_truncated_normal(self.a, self.b, self.mean, self.sigma, arr.shape)
+            _assignment(arr, tmp)
 
 
 def initializer(init, shape=None, dtype=mstype.float32):

mindspore/common/parameter.py

@@ -243,7 +243,8 @@ class Parameter(Tensor_):
     def __new__(cls, default_input, *args, **kwargs):
         init_data_flag = bool(isinstance(default_input, Tensor) and default_input.has_init)
         rc = sys.getrefcount(default_input)
-        input_class, *class_init_args = Parameter._get_parameter_new_args(default_input, rc)
+        init_param = getattr(cls, "init_param", True)
+        input_class, *class_init_args = Parameter._get_parameter_new_args(default_input, rc, init_param)
         new_type = Parameter._get_base_class(input_class)
         obj = input_class.__new__(new_type)
         input_class.__init__(obj, *class_init_args)
@@ -355,7 +356,7 @@ class Parameter(Tensor_):
         return new_type
 
     @staticmethod
-    def _get_parameter_new_args(data, rc):
+    def _get_parameter_new_args(data, rc, init_param=True):
         """Set `set_data` of current `Parameter`."""
         if isinstance(data, bool):
             raise ValueError('Parameter data can not be `bool`')
@@ -370,8 +371,8 @@ class Parameter(Tensor_):
                     return (Tensor, data.asnumpy(), mstype.qint4x2)
                 return (Tensor, data.asnumpy())
 
-            not_init_data = _is_role_sched() or (_is_role_pserver() and _cache_enable()
-                                                 ) or _is_in_auto_parallel_mode() or _is_parallel_mode()
+            not_init_data = not init_param or _is_role_sched() or (_is_role_pserver() and _cache_enable()) \
+                            or _is_in_auto_parallel_mode() or _is_parallel_mode()
             if not_init_data:
                 # do not init data while in auto parallel.
                 return (Tensor, None, data.dtype, get_slice_shape(data.dtype, data.shape), data.init)
@@ -976,7 +977,7 @@ class Parameter(Tensor_):
         """
         if self.is_default_input_init and self.is_in_parallel != _is_in_auto_parallel_mode():
             raise RuntimeError("Must set or change parallel mode before any initializer Tensor created.")
-        if self.init_mode is None:
+        if self.init_mode is None or not self.has_init:
             return self
         if self.inited_param is not None:
             return self.inited_param

mindspore/common/tensor.py

@@ -2896,8 +2896,13 @@ class Tensor(Tensor_, metaclass=_TensorMeta):
             self.slice_shape_of_persistent_data_ = data_shape
             self.slice_num_of_persistent_data_ = slice_num_of_persistent_data
 
+        from mindspore.common.initializer import Zero as ZeroInitializer
+
         try:
-            data = np.ndarray(data_shape, dtype=mstype.dtype_to_nptype(self.dtype))
+            if isinstance(self.init, ZeroInitializer):
+                data = np.zeros(data_shape, dtype=mstype.dtype_to_nptype(self.dtype))
+            else:
+                data = np.ndarray(data_shape, dtype=mstype.dtype_to_nptype(self.dtype))
         except ValueError as e:
             msg = "Error shape={}".format(shape)
             logger.critical(msg)
@@ -2933,7 +2938,7 @@ class Tensor(Tensor_, metaclass=_TensorMeta):
                     self.init.seed, _ = self.seed
 
         with seed_context(self.init):
-            if slice_num_of_persistent_data == 1:
+            if not isinstance(self.init, ZeroInitializer) and slice_num_of_persistent_data == 1:
                 self.init(data)
         self.init = None
 
@@ -4749,7 +4754,6 @@ class Tensor(Tensor_, metaclass=_TensorMeta):
         """
         return tensor_operator_registry.get('lu_solve')(self, LU_data, LU_pivots)
 
-
     def nextafter(self, other):
         r"""
         For details, please refer to :func:`mindspore.ops.nextafter`.
@@ -4763,7 +4767,6 @@ class Tensor(Tensor_, metaclass=_TensorMeta):
         validator.check_value_type('some', some, bool, 'Tensor.qr')
         return tensor_operator_registry.get('qr')(self, 'reduced' if some else 'complete')
 
-
     def ormqr(self, input2, input3, left=True, transpose=False):
         r"""
         For details, please refer to :func:`mindspore.ops.ormqr`,
@@ -4771,7 +4774,6 @@ class Tensor(Tensor_, metaclass=_TensorMeta):
         """
         return tensor_operator_registry.get('ormqr')(self, input2, input3, left, transpose)
 
-
     def masked_scatter(self, mask, x):
         r"""
         Returns a Tensor. Updates the value in the "self Tensor" with the `tensor` value according to the mask.
@@ -4812,7 +4814,6 @@ class Tensor(Tensor_, metaclass=_TensorMeta):
         """
         return tensor_operator_registry.get('masked_scatter')()(self, mask, x)
 
-
     def index_put(self, indices, values, accumulate=False):
         r"""
         Returns a Tensor. According to the index number of `indices` ,
@@ -4865,7 +4866,6 @@ class Tensor(Tensor_, metaclass=_TensorMeta):
         _index_put = tensor_operator_registry.get('index_put')(0 if accumulate is False else 1)
         return _index_put(self, values, indices)
 
-
     def move_to(self, to, blocking=True):
         r"""
         Copy Tensor to target device synchronously or asynchronously, default synchronously. only support PyNative mode.
@@ -4899,8 +4899,7 @@ class Tensor(Tensor_, metaclass=_TensorMeta):
         mode = context.get_context("mode")
         if mode != context.PYNATIVE_MODE:
             raise ValueError(f"The method of 'move_to' only supported in pynative mode, but got: {mode}.")
-        return Tensor(Tensor_.move_to(self, to, blocking), device="CPU" if to == "CPU" else None)
-
+        return Tensor_.move_to(self, to, blocking)
 
     def _offload(self):
         r"""
@@ -4946,44 +4945,6 @@ def _vm_compare(*args):
     return Tensor(np.array(fn(y)))
 
 
-def _check_sequence_shape(input_data):
-    """Check the shape of tensor input with type of sequence."""
-    max_dims_reached = False
-    max_ndim = 64 # corresponding to NPY_MAXDIMS
-    out_shape = [0]*max_ndim
-
-    def check_shape_recursive(input_data, curr_ndim):
-        nonlocal max_dims_reached, max_ndim, out_shape
-        if curr_ndim > max_ndim:
-            return False
-        if not isinstance(input_data, (tuple, list)):
-            if max_dims_reached and curr_ndim != max_ndim:
-                max_ndim = curr_ndim
-                return False
-            max_dims_reached = True
-            max_ndim = curr_ndim
-            return True
-        if not max_dims_reached:
-            out_shape[curr_ndim] = len(input_data)
-        else:
-            if out_shape[curr_ndim] != len(input_data):
-                max_ndim = curr_ndim
-                return False
-        if not input_data:
-            # process empty list
-            if not check_shape_recursive(None, curr_ndim + 1):
-                return False
-        for data in input_data:
-            if not check_shape_recursive(data, curr_ndim + 1):
-                return False
-        return True
-
-    if not check_shape_recursive(input_data, 0):
-        raise ValueError(f"When initializing a tensor with a sequence, the sequence has an inhomogeneous shape "
-                         f"after {max_ndim} dimensions. The detected shape was {tuple(out_shape[:max_ndim])} "
-                         f"+ inhomogeneous part.")
-
-
 def _check_tensor_input(input_data=None, dtype=None, shape=None, init=None):
     """Check the tensor input."""
     if input_data is not None and shape is not None:
@@ -4997,8 +4958,13 @@ def _check_tensor_input(input_data=None, dtype=None, shape=None, init=None):
         if isinstance(input_data, np.ndarray) and input_data.ndim >= 1 and input_data.size == 0:
             raise ValueError("input_data can not contain zero dimension.")
         if isinstance(input_data, (tuple, list)):
-            _check_sequence_shape(input_data)
-            if np.array(input_data).ndim >= 1 and np.array(input_data).size == 0:
+            try:
+                np_data = np.array(input_data)
+            except ValueError as e:
+                if "The requested array has an inhomogeneous shape" in str(e):
+                    raise TypeError(f"For Tensor, the input_data is {input_data} that contain unsupported element.")
+                raise
+            if np_data.ndim >= 1 and np_data.size == 0:
                 raise ValueError("input_data can not contain zero dimension.")
 
     if shape is not None and not (hasattr(init, "__enable_zero_dim__") and init.__enable_zero_dim__) and 0 in shape:

mindspore/communication/_comm_helper.py

@@ -499,3 +499,8 @@ def _destroy_group_helper(group):
         hccl.create_group(group)
     else:
         CollectiveManager.get_instance().destroy_group(group)
+
+
+def _get_group_map():
+    """Get the group map"""
+    return CollectiveManager.get_instance().get_group_map()

mindspore/communication/comm_func.py

@@ -209,7 +209,7 @@ def all_reduce(tensor, op=ReduceOp.SUM, group=GlobalComm.WORLD_COMM_GROUP, async
         RuntimeError: If device target is invalid, or backend is invalid, or distributed initialization fails.
 
     Supported Platforms:
-        ``Ascend`` ``GPU`` ``CPU``
+        ``Ascend``
 
     Examples:
         .. note::
@@ -275,7 +275,7 @@ def all_gather_into_tensor(tensor, group=GlobalComm.WORLD_COMM_GROUP, async_op=F
         RuntimeError: If device target is invalid, or backend is invalid, or distributed initialization fails.
 
     Supported Platforms:
-        ``Ascend`` ``GPU``
+        ``Ascend``
 
     Examples:
         .. note::
@@ -349,7 +349,7 @@ def reduce_scatter_tensor(tensor, op=ReduceOp.SUM, group=GlobalComm.WORLD_COMM_G
         RuntimeError: If device target is invalid, or backend is invalid, or distributed initialization fails.
 
     Supported Platforms:
-        ``Ascend`` ``GPU``
+        ``Ascend``
 
     Examples:
         .. note::
@@ -909,7 +909,7 @@ def send(tensor, dst=0, group=GlobalComm.WORLD_COMM_GROUP, tag=0):
         ValueError: If the rank ID of the process is greater than the rank size of the communication group.
 
     Supported Platforms:
-        ``Ascend`` ``GPU``
+        ``Ascend``
 
     Examples:
         .. note::
@@ -970,7 +970,7 @@ def recv(tensor, src=0, group=GlobalComm.WORLD_COMM_GROUP, tag=0):
         ValueError: If the rank ID of the process is greater than the rank size of the communication group.
 
     Supported Platforms:
-        ``Ascend`` ``GPU``
+        ``Ascend``
 
     Examples:
         .. note::
@@ -1040,7 +1040,7 @@ def isend(tensor, dst=0, group=GlobalComm.WORLD_COMM_GROUP, tag=0):
         ValueError: If the rank ID of the process is greater than the rank size of the communication group.
 
     Supported Platforms:
-        ``Ascend`` ``GPU``
+        ``Ascend``
 
     Examples:
         .. note::
@@ -1105,7 +1105,7 @@ def irecv(tensor, src=0, group=GlobalComm.WORLD_COMM_GROUP, tag=0):
         ValueError: If the rank ID of the process is greater than the rank size of the communication group.
 
     Supported Platforms:
-        ``Ascend`` ``GPU``
+        ``Ascend``
 
     Examples:
         .. note::

mindspore/context.py

@@ -936,6 +936,7 @@ def set_auto_parallel_context(**kwargs):
                \                 group_ckpt_save_file
                \                 auto_pipeline
                \                 dump_local_norm
+               \                 dump_device_local_norm
     ===========================  ===========================
 
     Args:
@@ -1090,6 +1091,9 @@ def set_auto_parallel_context(**kwargs):
         dump_local_norm (bool): Whether to dump local_norm value, when the `parallel_mode` is set to
                         ``semi_auto_parallel`` or ``auto_parallel``.
                         Default: ``False`` .
+        dump_device_local_norm (bool): Whether to dump device_local_norm value, when the `parallel_mode` is set to
+                        ``semi_auto_parallel`` or ``auto_parallel``.
+                        Default: ``False`` .
 
     Raises:
         ValueError: If input key is not attribute in auto parallel context.
@@ -1165,8 +1169,9 @@ def reset_auto_parallel_context():
     - pipeline_stages: 1.
     - pipeline_result_broadcast: False.
     - fusion_threshold: 64.
-    - dump_local_norm: False.
     - auto_pipeline: False.
+    - dump_local_norm: False.
+    - dump_device_local_norm: False.
 
     Examples:
         >>> import mindspore as ms
@@ -1686,6 +1691,15 @@ def set_context(**kwargs):
                 - 3: Optimize dataset reader with all scenes.
               - bias_add_comm_swap (bool): Enable node execution order swap communication operators and add operators
                 if ``True``. Only 1-dimension bias node is supported. Default: ``False``.
+              - enable_allreduce_slice_to_reducescatter (bool): Enable allreduce optimization. In the scenario where
+                the batchmatmul model introduces allreduce in parallel, if the subsequent nodes are stridedslice
+                operator with model parallel, allreduce will be optimized as reducescatter according to the identified
+                patterns. Typical used in MoE module with groupwise alltoall. Default: ``False``.
+              - enable_interleave_split_concat_branch (bool): Enable communication computation parallel optimization
+                for branches formed by split and concat operators with ``enable_interleave`` attribute. It is typical
+                used in MoE parallel scenario. After splitting the input data, each slice of data is processed by the
+                MoE module, and then the branch results are concatenated. When the optimization is enable,
+                communication and computation will be executed in parallel between branches. Default: ``False``.
             - host_scheduling_max_threshold(int): The max threshold to control whether the dynamic shape process is
               used when run the static graph, the default value is 0. When the number of operations in the static graph
               is less than the max threshold, this graph will be executed in dynamic shape process. In large model
@@ -1784,7 +1798,7 @@ def set_context(**kwargs):
             When both exist simultaneously, the global jit config will not overwrite the local network's jit config.
 
             - jit_level (str): Used to control the compilation optimization level. Default: ``""`` , The framework
-              automatically selects the execution method based on product, Altas training product is O2, and all other
+              automatically selects the execution method based on product, Atlas training product is O2, and all other
               products are O0. In addition, The option of the dynamic shape must be O0 or O1, O2 is not supported.
               The value range is as follows:
 

mindspore/dataset/engine/datasets_standard_format.py

@@ -33,6 +33,7 @@ from .datasets import UnionBaseDataset, SourceDataset, MappableDataset, Shuffle,
 from .datasets_user_defined import GeneratorDataset
 from .obs.obs_mindrecord_dataset import MindRecordFromOBS
 from .validators import check_csvdataset, check_minddataset, check_tfrecorddataset, check_obsminddataset
+from ..core.validator_helpers import type_check
 from ...mindrecord.config import _get_enc_key, _get_dec_mode, _get_hash_mode, decrypt, verify_file_hash
 
 
@@ -301,6 +302,22 @@ class MindDataset(MappableDataset, UnionBaseDataset):
                 else:
                     self.new_padded_sample[k] = v
 
+    def __deepcopy__(self, memodict):
+        if id(self) in memodict:
+            return memodict[id(self)]
+        return self.__safe_deepcopy__(memodict, exclude=("mindrecord_op"))
+
+    def __getitem__(self, index):
+        type_check(index, (int,), "index")
+        if index < 0:
+            raise ValueError("index cannot be negative, but got {0}.".format(index))
+        if not hasattr(self, "mindrecord_op"):
+            minddata_node = cde.MindDataNode(
+                self.dataset_files, self.columns_list, self.sampler, self.new_padded_sample,
+                self.num_padded, shuffle_to_shuffle_mode(self.shuffle_option))
+            self.mindrecord_op = minddata_node.Build()
+        return [t.as_array() for t in self.mindrecord_op[index]]
+
 
 class TFRecordDataset(SourceDataset, UnionBaseDataset):
     """

mindspore/dataset/engine/datasets_user_defined.py

@@ -19,6 +19,7 @@ After declaring the dataset object, you can further apply dataset operations
 (e.g. filter, skip, concat, map, batch) on it.
 """
 import builtins
+import copy
 import errno
 import itertools
 import math
@@ -50,6 +51,7 @@ from ..core.config import get_enable_shared_mem, get_prefetch_size, get_multipro
     get_enable_watchdog, get_debug_mode, get_seed, set_seed
 from ..core.datatypes import mstypelist_to_detypelist
 from ..core.py_util_helpers import ExceptionHandler
+from ..core.validator_helpers import type_check
 from ..transforms import transforms
 
 
@@ -427,6 +429,10 @@ class SamplerFn(cde.PythonMultiprocessingRuntime):
                 subprocess_file_descriptor = w.sentinel
                 st = time.time()
                 while _PythonMultiprocessing.is_process_alive(w.pid):
+                    process = psutil.Process(w.pid)
+                    if process.status() == psutil.STATUS_ZOMBIE:
+                        process.kill()
+                        break
                     time.sleep(0.01)  # sleep 10ms, waiting for the subprocess exit
                     if time.time() - st > check_interval:
                         logger.warning("Waiting for the subprocess worker [{}] to exit.".format(w.pid))
@@ -469,7 +475,7 @@ class SamplerFn(cde.PythonMultiprocessingRuntime):
 
                         # let the quit event notify the worker process to exit
                         w.join(timeout=5)
-                        if w.is_alive():
+                        if _PythonMultiprocessing.is_process_alive(w.pid):
                             # if the worker process did not exit, it may hang, try to terminate it
                             w.terminate()
                             w.close()
@@ -907,6 +913,26 @@ class GeneratorDataset(MappableDataset, UnionBaseDataset):
             return memodict[id(self)]
         return self.__safe_deepcopy__(memodict, exclude=("source", "__transfer_dataset__"))
 
+    def __getitem__(self, index):
+        type_check(index, (int, np.number), "index")
+        if not hasattr(self.source, "__getitem__"):
+            raise RuntimeError("Dataset don't support randomized access.")
+        if not hasattr(self, "generator_op"):
+            dataset = copy.deepcopy(self)
+            self.prepared_source = _generator_fn_wrapper(_cpp_sampler_fn, self.source)
+            if self.schema is None:
+                dataset.generator_node = cde.GeneratorNode(self.prepared_source, self.column_names, self.column_types,
+                                                           self.source_len, self.sampler, 1, None)
+            else:
+                schema = self.schema
+                if isinstance(schema, Schema):
+                    schema = self.schema.cpp_schema
+                dataset.generator_node = cde.GeneratorNode(self.prepared_source, schema, self.source_len,
+                                                           self.sampler, 1, None)
+            self.generator_op = dataset.generator_node.Build()
+        sample_id = self.generator_op.GetMappedIndex(index)
+        return self.source[sample_id]
+
     def is_shuffled(self):
         if self.sampler:
             return self.sampler.is_shuffled()

mindspore/experimental/llm_boost/__init__.py

@@ -15,7 +15,7 @@
 """LlmBoost Register"""
 from __future__ import absolute_import
 
-from mindspore.experimental.llm_boost.atb import *
+from mindspore.experimental.llm_boost.atb import LlamaBoost, QwenBoost
 from mindspore.experimental.llm_boost.register import LlmBoostRegister
 
-__all__ = ['LlmBoostRegister']
+__all__ = ["LlmBoostRegister"]