# JAXとの併用

このドキュメントは、`datasets`をJAXと組み合わせて使用する際の簡単な入門ガイドです。特に、データセットから`jax.Array`オブジェクトを取得する方法、およびそれらを使用してJAXモデルを学習させる方法に焦点を当てています。

> [!TIP]
> 上記のコードを再現するには`jax`と`jaxlib`が必要です。`pip install datasets[jax]`としてインストールしてください。

## データセットの形式

デフォルトでは、datasetsは通常のPythonオブジェクト（整数、浮動小数点数、文字列、リストなど）を返します。JAXは数値のみをサポートするため、文字列およびバイナリオブジェクトは変更されません。

代わりにJAX配列（numpy類似）を取得するには、データセットの形式を`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]
> [Dataset](/docs/datasets/v5.2.0/ja/package_reference/main_classes#datasets.Dataset) オブジェクトは Arrow テーブルのラッパーであり、データセット内の配列から JAX 配列への高速な読み取りを可能にします。

`DatasetDict` オブジェクトにもまったく同じ手順が適用されることに注意してください。`DatasetDict` のフォーマットを `jax` に設定すると、そこに含まれるすべての `Dataset` が `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)}
```

もう一つ考慮すべき点は、フォーマットが実際にデータにアクセスするまで適用されないことです。したがって、データセットからJAX配列を取得したい場合は、まずデータにアクセスする必要があります。そうしないと、フォーマットは変更されません。

最後に、希望するデバイスにデータをロードするには、`device`引数を指定できますが、`jaxlib.xla_extension.Device`は`pickle`でも`dill`でもシリアライズできないためサポートされていません。そのため、代わりにその文字列識別子を使用する必要があります:

```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
```

`device`引数が`with_format`に指定されない場合、デフォルトのデバイスである`jax.devices()[0]`が使用されることに注意してください。

### N次元配列

データセットがN次元の配列で構成されている場合、デフォルトでは、shapeが固定されている場合、それらは同じtensorとして扱われることに注意してください。

```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)]}
```

ただし、このロジックは多くの場合、遅い形状の比較とデータのコピーを必要とします。これを回避するには、`Array`機能型を明示的に使用し、テンソルの形状を指定する必要があります。

```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)}
```

### その他の特徴タイプ

[ClassLabel](/docs/datasets/v5.2.0/ja/package_reference/main_classes#datasets.ClassLabel) データは正しく配列に変換されます:

```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)}
```

JAXは数値のみをサポートするため、文字列オブジェクトとバイナリオブジェクトは変更されません。

[Image](/docs/datasets/v5.2.0/ja/package_reference/main_classes#datasets.Image)および[Audio](/docs/datasets/v5.2.0/ja/package_reference/main_classes#datasets.Audio)の機能タイプもサポートされています。

> [!TIP]
> [Image](/docs/datasets/v5.2.0/ja/package_reference/main_classes#datasets.Image)機能タイプを使用するには、`vision`エクストラを`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]
> [Audio](/docs/datasets/v5.2.0/ja/package_reference/main_classes#datasets.Audio)機能タイプを使用するには、`audio`エクストラを`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)
```

## データ読み込み

JAXには組み込みのデータ読み込み機能がないため、[PyTorch](https://pytorch.org/)の`DataLoader`または[TensorFlow](https://www.tensorflow.org/)の`tf.data.Dataset`などのライブラリを使用してデータをロードする必要があります。この点について、[JAXのドキュメント](https://jax.readthedocs.io/en/latest/notebooks/Neural_Network_and_Data_Loading.html#data-loading-with-pytorch)には次のように記載されています。「JAXはプログラムの変換とアクセラレータ対応のNumPyに特化しているため、JAXライブラリにはデータ読み込みやデータ整形は含まれていません。優れたデータローダーはすでに数多く存在しているため、何かを再発明するのではなく、それらを利用しましょう。ここではPyTorchのデータローダーを取得し、NumPy配列で動作させるための小さなシムを作成します。」

`datasets`におけるJAX形式のデータが非常に有用である理由は、データ読み込みの部分を気にすることなく、HuggingFace Hubの任意のモデルをJAXで使用できるためです。

### `with_format('jax')`の使用

データセットからJAX配列を取得する最も簡単な方法は、`with_format('jax')`メソッドを使用することです。HuggingFace Hubのhttps://huggingface.co/datasets/ylecun/mnistで利用可能な[MNISTデータセット](http://yann.lecun.com/exdb/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)}
```

形式が設定されたら、`Dataset.iter()` メソッドを使用して、JAX モデルにバッチ単位でデータセットを供給できます:

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