# Sử dụng với JAX

Tài liệu này là phần giới thiệu nhanh về cách sử dụng `datasets` với JAX, tập trung đặc biệt vào cách lấy các đối tượng `jax.Array` từ các dataset của chúng ta và cách sử dụng chúng để huấn luyện các mô hình JAX.

> [!TIP]
> `jax` và `jaxlib` là bắt buộc để tái tạo mã ở trên, vì vậy hãy đảm bảo bạn cài đặt chúng bằng `pip install datasets[jax]`.

## Định dạng tập dữ liệu

Mặc định, các tập dữ liệu trả về các đối tượng Python thông thường: số nguyên, số thực, chuỗi, danh sách, v.v., và các đối tượng chuỗi và nhị phân không bị thay đổi, vì JAX chỉ hỗ trợ các con số.

Để nhận các mảng JAX (tương tự numpy) thay vào đó, bạn có thể đặt định dạng của tập dữ liệu thành `jax`:

```py
>>> from datasets import Dataset
>>> data = [[1, 2], [3, 4]]
>>> ds = Dataset.from_dict({"data": data})
>>> ds = ds.with_format("jax")
>>> ds[0]
{'data': DeviceArray([1, 2], dtype=int32)}
>>> ds[:2]
{'data': DeviceArray([
    [1, 2],
    [3, 4]], dtype=int32)}
```

> [!TIP]
> Một đối tượng [Dataset](/docs/datasets/main/vi/package_reference/main_classes#datasets.Dataset) là lớp bọc của một bảng Arrow, cho phép đọc nhanh từ các mảng trong bộ dữ liệu sang các mảng JAX.

Lưu ý rằng quy trình hoàn toàn tương tự cũng áp dụng cho các đối tượng `DatasetDict`, do đó khi đặt định dạng của một `DatasetDict` thành `jax`, tất cả các `Dataset` bên trong sẽ được định dạng thành `jax`:

```py
>>> from datasets import DatasetDict
>>> data = {"train": {"data": [[1, 2], [3, 4]]}, "test": {"data": [[5, 6], [7, 8]]}}
>>> dds = DatasetDict.from_dict(data)
>>> dds = dds.with_format("jax")
>>> dds["train"][:2]
{'data': DeviceArray([
    [1, 2],
    [3, 4]], dtype=int32)}
```

Một điều khác bạn cần lưu ý là định dạng chỉ được áp dụng khi bạn thực sự truy cập dữ liệu. Vì vậy, nếu bạn muốn lấy một mảng JAX từ một dataset, bạn cần truy cập dữ liệu trước, nếu không định dạng sẽ không thay đổi.

Cuối cùng, để nạp dữ liệu vào thiết bị bạn mong muốn, bạn có thể chỉ định đối số `device`, nhưng lưu ý rằng `jaxlib.xla_extension.Device` không được hỗ trợ vì nó không thể được serialize bằng cả `pickle` lẫn `dill`, do đó bạn sẽ cần sử dụng định danh dạng chuỗi của nó thay thế:

```py
>>> import jax
>>> from datasets import Dataset
>>> data = [[1, 2], [3, 4]]
>>> ds = Dataset.from_dict({"data": data})
>>> device = str(jax.devices()[0])  # Not casting to `str` before passing it to `with_format` will raise a `ValueError`
>>> ds = ds.with_format("jax", device=device)
>>> ds[0]
{'data': DeviceArray([1, 2], dtype=int32)}
>>> ds[0]["data"].device()
TFRT_CPU_0
>>> assert ds[0]["data"].device() == jax.devices()[0]
True
```

Lưu ý rằng nếu không cung cấp đối số `device` cho `with_format` thì nó sẽ sử dụng thiết bị mặc định là `jax.devices()[0]`.

### Mảng N chiều

Nếu tập dữ liệu của bạn bao gồm các mảng N chiều, bạn sẽ thấy rằng mặc định chúng được coi là cùng một tensor nếu hình dạng được cố định:

```py
>>> from datasets import Dataset
>>> data = [[[1, 2],[3, 4]], [[5, 6],[7, 8]]]  # fixed shape
>>> ds = Dataset.from_dict({"data": data})
>>> ds = ds.with_format("jax")
>>> ds[0]
{'data': Array([[1, 2],
        [3, 4]], dtype=int32)}
```

```py
>>> from datasets import Dataset
>>> data = [[[1, 2],[3]], [[4, 5, 6],[7, 8]]]  # varying shape
>>> ds = Dataset.from_dict({"data": data})
>>> ds = ds.with_format("jax")
>>> ds[0]
{'data': [Array([1, 2], dtype=int32), Array([3], dtype=int32)]}
```

Tuy nhiên, logic này thường đòi hỏi các phép so sánh shape và sao chép dữ liệu chậm. Để tránh điều này, bạn phải sử dụng rõ ràng kiểu đặc trưng `Array` và chỉ định shape của các tensor của bạn:

```py
>>> from datasets import Dataset, Features, Array2D
>>> data = [[[1, 2],[3, 4]],[[5, 6],[7, 8]]]
>>> features = Features({"data": Array2D(shape=(2, 2), dtype='int32')})
>>> ds = Dataset.from_dict({"data": data}, features=features)
>>> ds = ds.with_format("jax")
>>> ds[0]
{'data': Array([[1, 2],
        [3, 4]], dtype=int32)}
>>> ds[:2]
{'data': Array([[[1, 2],
         [3, 4]],
 
        [[5, 6],
         [7, 8]]], dtype=int32)}
```

### Các loại tính năng khác

Dữ liệu [ClassLabel](/docs/datasets/main/vi/package_reference/main_classes#datasets.ClassLabel) được chuyển đổi thành mảng một cách chính xác:

```py
>>> from datasets import Dataset, Features, ClassLabel
>>> labels = [0, 0, 1]
>>> features = Features({"label": ClassLabel(names=["negative", "positive"])})
>>> ds = Dataset.from_dict({"label": labels}, features=features)
>>> ds = ds.with_format("jax")
>>> ds[:3]
{'label': DeviceArray([0, 0, 1], dtype=int32)}
```

Các đối tượng chuỗi và nhị phân không thay đổi, vì JAX chỉ hỗ trợ các số.

Các loại tính năng [Image](/docs/datasets/main/vi/package_reference/main_classes#datasets.Image) và [Audio](/docs/datasets/main/vi/package_reference/main_classes#datasets.Audio) cũng được hỗ trợ.

> [!TIP]
> Để sử dụng kiểu tính năng [Image](/docs/datasets/main/vi/package_reference/main_classes#datasets.Image), bạn cần cài đặt phần mở rộng `vision` bằng lệnh `pip install datasets[vision]`.

```py
>>> from datasets import Dataset, Features, Image
>>> images = ["path/to/image.png"] * 10
>>> features = Features({"image": Image()})
>>> ds = Dataset.from_dict({"image": images}, features=features)
>>> ds = ds.with_format("jax")
>>> ds[0]["image"].shape
(512, 512, 3)
>>> ds[0]
{'image': DeviceArray([[[ 255, 255, 255],
              [ 255, 255, 255],
              ...,
              [ 255, 255, 255],
              [ 255, 255, 255]]], dtype=uint8)}
>>> ds[:2]["image"].shape
(2, 512, 512, 3)
>>> ds[:2]
{'image': DeviceArray([[[[ 255, 255, 255],
              [ 255, 255, 255],
              ...,
              [ 255, 255, 255],
              [ 255, 255, 255]]]], dtype=uint8)}
```

> [!TIP]
> Để sử dụng kiểu tính năng [Audio](/docs/datasets/main/vi/package_reference/main_classes#datasets.Audio), bạn cần cài đặt phần mở rộng `audio` bằng lệnh `pip install datasets[audio]`.

```py
>>> from datasets import Dataset, Features, Audio
>>> audio = ["path/to/audio.wav"] * 10
>>> features = Features({"audio": Audio()})
>>> ds = Dataset.from_dict({"audio": audio}, features=features)
>>> ds = ds.with_format("jax")
>>> ds[0]["audio"]["array"]
DeviceArray([-0.059021  , -0.03894043, -0.00735474, ...,  0.0133667 ,
              0.01809692,  0.00268555], dtype=float32)
>>> ds[0]["audio"]["sampling_rate"]
DeviceArray(44100, dtype=int32, weak_type=True)
```

## Tải dữ liệu

JAX không có khả năng tải dữ liệu tích hợp sẵn, vì vậy bạn sẽ cần sử dụng một thư viện như [PyTorch](https://pytorch.org/) để tải dữ liệu của bạn bằng `DataLoader` hoặc [TensorFlow](https://www.tensorflow.org/) bằng `tf.data.Dataset`. Trích dẫn từ [tài liệu JAX](https://jax.readthedocs.io/en/latest/notebooks/Neural_Network_and_Data_Loading.html#data-loading-with-pytorch) về chủ đề này: "JAX tập trung hoàn toàn vào việc chuyển đổi chương trình và NumPy có hỗ trợ từ bộ tăng tốc, vì vậy chúng tôi không bao gồm việc tải hoặc xử lý dữ liệu trong thư viện JAX. Đã có rất nhiều trình tải dữ liệu tuyệt vời ngoài kia, vì vậy hãy chỉ sử dụng chúng thay vì phát minh lại bất cứ thứ gì. Chúng tôi sẽ lấy trình tải dữ liệu của PyTorch và tạo một shim nhỏ để nó hoạt động với các mảng NumPy."

Đó là lý do tại sao việc định dạng JAX trong `datasets` lại hữu ích đến vậy, vì nó cho phép bạn sử dụng bất kỳ mô hình nào từ HuggingFace Hub với JAX, mà không cần phải lo lắng về phần tải dữ liệu.

### Sử dụng `with_format('jax')`

Cách dễ nhất để lấy các mảng JAX từ một dataset là sử dụng phương thức `with_format('jax')`. Giả sử chúng ta muốn huấn luyện một mạng nơ-ron trên [dataset MNIST](http://yann.lecun.com/exdb/mnist/) có sẵn trên HuggingFace Hub tại https://huggingface.co/datasets/ylecun/mnist.

```py
>>> from datasets import load_dataset
>>> ds = load_dataset("ylecun/mnist")
>>> ds = ds.with_format("jax")
>>> ds["train"][0]
{'image': DeviceArray([[  0,   0,   0, ...],
                       [  0,   0,   0, ...],
                       ...,
                       [  0,   0,   0, ...],
                       [  0,   0,   0, ...]], dtype=uint8),
 'label': DeviceArray(5, dtype=int32)}
```

Khi định dạng đã được thiết lập, chúng ta có thể cung cấp bộ dữ liệu cho mô hình JAX theo từng lô bằng cách sử dụng phương thức `Dataset.iter()`:

```py
>>> for epoch in range(epochs):
...     for batch in ds["train"].iter(batch_size=32):
...         x, y = batch["image"], batch["label"]
...         ...
```