code: def prepare_jittor_dataloader(
      ds_or_db,
      batch_size: int = 16,
      shuffle: bool = None,
      drop_last: bool = False,
      num_workers: int = 0,
      buffer_size: int = 512 * 1024 * 1024,
      stop_grad: bool = True,
      keep_numpy_array: bool = False,
      endless: bool = False,
      collate_fn: Union[None, str, Callable] = "auto",
      non_train_batch_size: int = None,
  ) -> Union[Dict[str, JittorDataLoader], JittorDataLoader]:
      """
      ``prepare_jittor_dataloader`` 的功能是将输入的单个或多个 dataset 同时转为 :class:`JittorDataLoader` 对象， 详见 :class:`~fastNLP.core.dataloaders.JittorDataLoader`。
      根据 ``ds_or_db`` 的类型 ``[DataSet, DataBundle, Dict[name, Dataset]]`` 不同而有不同返回结果, 具体如下:
  
          * 当 ds_or_db 为 :class:`~fastNLP.io.DataSet` 时，``prepare_jittor_dataloader`` 会将使用的除了 non_train_batch_size 和 non_train_sampler 以外的参数来
            帮你实例化一个 :class:`JittorDataLoader` 对象并返回该对象。 详见 :class:`~fastNLP.core.dataloaders.JittorDataLoader`；
          * 当 ds_or_db 为 :class:`~fastNLP.io.DataBundle` 时，``prepare_jittor_dataloader`` 会遍历 ``DataBundle`` 的数据集的 key-value
            来创建不同的 :class:`JittorDataLoader` 对象；当 key 中包含 ``'train'`` 字符串时，``prepare_jittor_dataloader`` 默认该 value 为训练数据集，
            会将 ``batch_size`` 和 ``sampler`` 作为参数，其他 key 不包含 ``'train'`` 字符串的数据集则使用 ``non_train_size`` 和 ``non_train_sampler`` 作为参数。
            最终根据 ``key: JittorDataLoader`` 组成 ``Dict[key, JittorDataLoader]`` 的字典返回；
          * 当 ds_or_db 为 ``Dict[str, DataSet]`` 字典类型时， ``prepare_jittor_dataloader`` 会遍历 该 dict 的的 key-value 来创建不同的
            :class:`JittorDataLoader` 对象；当 key 中包含 ``'train'`` 字符串时，``prepare_jittor_dataloader`` 默认该 value 为训练数据集，会将 ``batch_size`` 和
            ``sampler`` 作为参数，其他 key 不包含 ``'train'`` 字符串的数据集则使用 ``non_train_size`` 和 ``non_train_sampler`` 作为参数。最终根据  ``key: JittorDataLoader`` 组成
            ``Dict[key, JittorDataLoader]`` 的字典返回；
  
      :param ds_or_db: 可以有以下三种取值：
  
          * ds_or_db 为 :class:`~fastNLP.io.DataBundle`, 返回值为 ``Dict[str, TorchDataLoader]`` 的字典；
          * ds_or_db 为 ``Dict[str, DataSet]`` 字典， 返回值为 ``Dict[str, TorchDataLoader]`` 的字典；
          * ds_or_db 为实现了 :meth:`__getitem__` 和 :meth:`__len__` 的对象 ，返回值为 :class:`~fastNLP.core.dataloaders.JittorDataLoader`；
  
      :param non_train_batch_size: 如果传入的 ``ds_or_db`` 为 :class:`Dict` 或 :class:`~fastNLP.io.DataBundle` 对象，可以通过改参数
          设置名称不为 `train` 的其他 ``dataset`` 的 ``batch_size``。 默认为 ``16``。
      :param batch_size: 批次大小，默认为 ``16`` 且当 batch_sampler 为 None 有效。
      :param shuffle: 是否打乱数据集， 默认为 ``None``, 如果传入的 ``ds_or_db`` 可以判断出哪个是 ``'train'`` 则设置其 shuffle 为 True ，
          其它的为 False 。
      :param drop_last: 当 ``drop_last=True`` 时，:class:`JittorDataLoader` 会扔掉最后一个长度小于 ``batch_size`` 的 batch 数据;
          若 ``drop_last=False`` , 则会返回该 batch 数据。 默认为 ``False`` 。
      :param num_workers: 当 ``num_workers > 0`` 时, :class:`JittorDataLoader` 会开启 num_workers 个子进程来处理数据， 可以加快
          数据处理速度，但同时也消耗大量内存。 当 ``num_workers=0`` 时， 不开启子进程。 默认为 ``0``。
      :param buffer_size: 每个进程占用的内存空间，默认为512M。主要是配合 ``num_workers`` 使用，用户可以自定义每个进程的内存大小。
      :param stop_grad: 是否不使用梯度， 默认 ``True`` 。
      :param keep_numpy_array: 返回的数据是 :class:`np.array` 类型而不是 :class:`ittor.Var` 类型，默认为 ``False``
      :param endless: 是否让 :class:`JittorDataLoader` 无限返回数据，也就是将 dataset 循环使用使得返回数据是没有限制的。默认为 ``False``.
      :param collate_fn: 用于从 dataset 取到的一个 batch 数据进行打包处理的 Callable 函数，其值应该为以下三个: ``[None, "auto", Callable]``.
  
          * callate_fn 为 ``None`` 时，需要注意的是此时传进来的 datset 类型不能为 :class:`~fastNLP.core.dataset.DataSet` , 当 collate_fn 为 ``None`` 时，
            :class:`JittorDataLoader` 调用默认的 Jittor 框架的 ``DataLoader`` 自带的 ``collate_batch`` 作为 callate_fn 的默认值， 其无法处理
            :class:`~fastNLP.core.dataset.DataSet` 的 dataset 对象。
          * callate_fn 为 ``'auto'`` 时，:class:`JittorDataLoader` 使用 :class:`~fastNLP.core.collators.Collator` 作为 collate_fn 的默认值。
            此时可以配套使用 :class:`JittorDataLoader` 的 ``set_pad`` 和 ``set_ignore`` 方法来设置 pad_val 或 忽略某个 field 的检测。
          * collate_fn 为 :class:`Callable` 时， 该 Callable 函数应当接受一个 batch 参数作为输入， batch 是一个 List 对象且 List 中的每一条数据都是
            dataset 的一条数据；该 Callable 函数还应当返回一个对象。
      """
      from fastNLP.io.data_bundle import DataBundle
  
      if isinstance(ds_or_db, DataBundle):
          dl_bundle = {}
          for name, ds in ds_or_db.iter_datasets():
              if "train" in name:
                  dl_bundle[name] = JittorDataLoader(
                      ds,
                      batch_size=batch_size,
                      shuffle=True if shuffle is None else shuffle,
                      drop_last=drop_last,
                      num_workers=num_workers,
                      buffer_size=buffer_size,
                      stop_grad=stop_grad,
                      keep_numpy_array=keep_numpy_array,
                      endless=endless,
                      collate_fn=collate_fn,
                  )
              else:
                  dl_bundle[name] = JittorDataLoader(
                      ds,
                      batch_size=non_train_batch_size
                      if non_train_batch_size
                      else batch_size,
                      shuffle=False if shuffle is None else shuffle,
                      drop_last=drop_last,
                      num_workers=num_workers,
                      buffer_size=buffer_size,
                      stop_grad=stop_grad,
                      keep_numpy_array=keep_numpy_array,
                      endless=endless,
                      collate_fn=collate_fn,
                  )
          return dl_bundle
  
      elif isinstance(ds_or_db, Dict):
          ds_dict = {}
          for name, ds in ds_or_db.items():
              if "train" in name:
                  dl = JittorDataLoader(
                      ds,
                      batch_size=batch_size,
                      shuffle=True if shuffle is None else shuffle,
                      drop_last=drop_last,
                      num_workers=num_workers,
                      buffer_size=buffer_size,
                      stop_grad=stop_grad,
                      keep_numpy_array=keep_numpy_array,
                      endless=endless,
                      collate_fn=collate_fn,
                  )
              else:
                  dl = JittorDataLoader(
                      ds,
                      batch_size=non_train_batch_size
                      if non_train_batch_size
                      else batch_size,
                      shuffle=False if shuffle is None else shuffle,
                      drop_last=drop_last,
                      num_workers=num_workers,
                      buffer_size=buffer_size,
                      stop_grad=stop_grad,
                      keep_numpy_array=keep_numpy_array,
                      endless=endless,
                      collate_fn=collate_fn,
                  )
              ds_dict[name] = dl
          return ds_dict
  
      elif isinstance(ds_or_db, HasLenGetitemType):
          dl = JittorDataLoader(
              ds_or_db,
              batch_size=batch_size,
              shuffle=False if shuffle is None else shuffle,
              drop_last=drop_last,
              num_workers=num_workers,
              buffer_size=buffer_size,
              stop_grad=stop_grad,
              keep_numpy_array=keep_numpy_array,
              endless=endless,
              collate_fn=collate_fn,
          )
          return dl
  
      else:
          raise ValueError(
              f"ds_or_db: {ds_or_db} must be fastnlp dataset or data_bundle or mapping!"
          )
  
imports: from typing import Callable, Optional, List, Union, Dict, Sequence
  from copy import deepcopy
  import numpy
  from fastNLP.envs.imports import _NEED_IMPORT_JITTOR
  from fastNLP.core.collators import Collator
  from fastNLP.core.dataloaders.utils import indice_collate_wrapper
  from fastNLP.core.dataset import DataSet
  from ..utils import HasLenGetitemType
label_desc: Use image node_id 1 for calling a global function with example usage: prepare_jittor_dataloader(ds_or_db, batch_size, shuffle, drop_last, num_workers, buffer_size, stop_grad, keep_numpy_array, endless, collate_fn, non_train_batch_size) and returns: dl_bundle, ds_dict, dl