cyreal 0.1.1__tar.gz

cyreal-0.1.1/PKG-INFO

@@ -0,0 +1,179 @@
+Metadata-Version: 2.4
+Name: cyreal
+Version: 0.1.1
+Summary: Jittable data loading utilities for JAX.
+Author: 
+Classifier: Programming Language :: Python :: 3
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Operating System :: OS Independent
+Requires-Python: >=3.10
+Description-Content-Type: text/markdown
+Requires-Dist: jax
+Requires-Dist: jaxlib
+Requires-Dist: numpy
+Provides-Extra: dev
+Requires-Dist: pytest; extra == "dev"
+Dynamic: classifier
+Dynamic: description
+Dynamic: description-content-type
+Dynamic: provides-extra
+Dynamic: requires-dist
+Dynamic: requires-python
+Dynamic: summary
+
+# Cyreal - Another JAX DataLoader 
+
+> `grain` for the corporations, `cyreal` for the people
+
+Pure `jax` utilities for iterating over finite datasets without ever touching `torch` or `tensorflow`. Dataloaders support `jax.jit`, `jax.grad`, `jax.lax.scan`, and other function transformations.
+
+## Installation
+
+The only dependency is `jax`. On GPU machines, install the
+appropriate JAX build for your CUDA version.
+
+`pip install cyreal`
+
+## Quick start with MNIST
+Write `torch`-style dataloaders without `torch`
+
+```python
+import jax
+import jax.numpy as jnp
+
+from cyreal import (
+  ArraySampleSource,
+  BatchTransform,
+  DataLoader,
+  DevicePutTransform,
+  MNISTDataset,
+)
+
+train_data = MNISTDataset(split="train").as_array_dict()
+pipeline = [
+  ArraySampleSource(train_data, ordering="shuffle"),
+  BatchTransform(batch_size=128),
+  DevicePutTransform(),
+]
+loader = DataLoader(pipeline=pipeline)
+state = loader.init_state(jax.random.PRNGKey(0))
+
+for batch, mask in loader.iterate(state):
+  ...  # train your network!
+```
+
+## Scan and Avoid Boilerplate 
+
+`DataLoader.scan_epoch` will run a full pass through the dataset into a single
+`jax.lax.scan` to minimize dispatch overhead. This will `jit` the `body_fn`.
+
+```python
+def body_fn(model_state, batch, mask):
+  model_state = update_model(model_state, batch, mask)
+  return model_state, None
+
+loader_state, model_state, _ = loader.scan_epoch(loader_state, model_state, body_fn)
+```
+
+## JIT Capabilities
+Do you enjoy premature optimization? Why not `jit` the entire train epoch?
+
+
+```python
+import jax
+import jax.numpy as jnp
+
+from cyreal import (
+  ArraySampleSource,
+  BatchTransform,
+  DataLoader,
+  DevicePutTransform,
+  MNISTDataset,
+)
+
+train_data = MNISTDataset(split="train").as_array_dict()
+pipeline = [
+  ArraySampleSource(train_data, ordering="shuffle"),
+  BatchTransform(batch_size=128),
+  DevicePutTransform(),
+]
+loader = DataLoader(pipeline)
+loader_state = loader.init_state(jax.random.PRNGKey(0))
+model_state = model_init()
+
+@jax.jit
+def train_epoch(model_state, loader_state):
+  def body_fn(model_state, batch, mask):
+    # Update the network using your train fn
+    new_model_state = model_update(model_state, batch, mask)
+    return new_model_state, None
+
+  loader_state, model_state, _ = loader.scan_epoch(loader_state, model_state, body_fn)
+  return model_state, loader_state
+
+model_state, loader_state = train_epoch(model_state, loader_state)
+```
+
+
+## Streaming from Disk
+Is your dataset enormous? Swap in a disk-backed source.
+
+```python
+import jax
+
+from cyreal import (
+  BatchTransform,
+  DataLoader,
+  DevicePutTransform,
+  MNISTDataset,
+)
+
+pipeline = [
+  MNISTDataset.make_disk_source(split="train", ordering="shuffle", prefetch_size=1024),
+  BatchTransform(batch_size=128),
+  DevicePutTransform(),
+]
+
+loader = DataLoader(pipeline=pipeline)
+state = loader.init_state(jax.random.PRNGKey(0))
+
+for batch, mask in loader.iterate(state):
+  ...  # stream without holding the dataset in RAM
+```
+
+## For the Dirty and Impure
+Want to `jit` but also log some metrics? Use `HostCallbackTransform` which utilizes `jax.experimental.io_callback` under the hood.
+
+```python
+import jax.numpy as jnp
+import numpy as np
+
+from cyreal import (
+  ArraySampleSource,
+  BatchTransform,
+  DataLoader,
+  HostCallbackTransform,
+  MNISTDataset,
+)
+
+def model(images):
+  return jnp.mean(images.astype(jnp.float32), axis=(1, 2, 3))
+
+def cross_entropy(logits, labels):
+  labels = labels.astype(jnp.float32)
+  return (logits - labels) ** 2
+
+def log_loss(batch, mask):
+  logits = model(batch["image"])
+  loss = jnp.mean(cross_entropy(logits, batch["label"]) * mask[:, None])
+  print("loss:", float(np.asarray(loss)))
+  return batch
+
+loader = DataLoader(
+  pipeline=[
+    ArraySampleSource(MNISTDataset(split="train").as_array_dict(), ordering="shuffle"),
+    BatchTransform(batch_size=128),
+    HostCallbackTransform(fn=log_loss),
+  ],
+)
+```

cyreal-0.1.1/README.md

@@ -0,0 +1,156 @@
+# Cyreal - Another JAX DataLoader 
+
+> `grain` for the corporations, `cyreal` for the people
+
+Pure `jax` utilities for iterating over finite datasets without ever touching `torch` or `tensorflow`. Dataloaders support `jax.jit`, `jax.grad`, `jax.lax.scan`, and other function transformations.
+
+## Installation
+
+The only dependency is `jax`. On GPU machines, install the
+appropriate JAX build for your CUDA version.
+
+`pip install cyreal`
+
+## Quick start with MNIST
+Write `torch`-style dataloaders without `torch`
+
+```python
+import jax
+import jax.numpy as jnp
+
+from cyreal import (
+  ArraySampleSource,
+  BatchTransform,
+  DataLoader,
+  DevicePutTransform,
+  MNISTDataset,
+)
+
+train_data = MNISTDataset(split="train").as_array_dict()
+pipeline = [
+  ArraySampleSource(train_data, ordering="shuffle"),
+  BatchTransform(batch_size=128),
+  DevicePutTransform(),
+]
+loader = DataLoader(pipeline=pipeline)
+state = loader.init_state(jax.random.PRNGKey(0))
+
+for batch, mask in loader.iterate(state):
+  ...  # train your network!
+```
+
+## Scan and Avoid Boilerplate 
+
+`DataLoader.scan_epoch` will run a full pass through the dataset into a single
+`jax.lax.scan` to minimize dispatch overhead. This will `jit` the `body_fn`.
+
+```python
+def body_fn(model_state, batch, mask):
+  model_state = update_model(model_state, batch, mask)
+  return model_state, None
+
+loader_state, model_state, _ = loader.scan_epoch(loader_state, model_state, body_fn)
+```
+
+## JIT Capabilities
+Do you enjoy premature optimization? Why not `jit` the entire train epoch?
+
+
+```python
+import jax
+import jax.numpy as jnp
+
+from cyreal import (
+  ArraySampleSource,
+  BatchTransform,
+  DataLoader,
+  DevicePutTransform,
+  MNISTDataset,
+)
+
+train_data = MNISTDataset(split="train").as_array_dict()
+pipeline = [
+  ArraySampleSource(train_data, ordering="shuffle"),
+  BatchTransform(batch_size=128),
+  DevicePutTransform(),
+]
+loader = DataLoader(pipeline)
+loader_state = loader.init_state(jax.random.PRNGKey(0))
+model_state = model_init()
+
+@jax.jit
+def train_epoch(model_state, loader_state):
+  def body_fn(model_state, batch, mask):
+    # Update the network using your train fn
+    new_model_state = model_update(model_state, batch, mask)
+    return new_model_state, None
+
+  loader_state, model_state, _ = loader.scan_epoch(loader_state, model_state, body_fn)
+  return model_state, loader_state
+
+model_state, loader_state = train_epoch(model_state, loader_state)
+```
+
+
+## Streaming from Disk
+Is your dataset enormous? Swap in a disk-backed source.
+
+```python
+import jax
+
+from cyreal import (
+  BatchTransform,
+  DataLoader,
+  DevicePutTransform,
+  MNISTDataset,
+)
+
+pipeline = [
+  MNISTDataset.make_disk_source(split="train", ordering="shuffle", prefetch_size=1024),
+  BatchTransform(batch_size=128),
+  DevicePutTransform(),
+]
+
+loader = DataLoader(pipeline=pipeline)
+state = loader.init_state(jax.random.PRNGKey(0))
+
+for batch, mask in loader.iterate(state):
+  ...  # stream without holding the dataset in RAM
+```
+
+## For the Dirty and Impure
+Want to `jit` but also log some metrics? Use `HostCallbackTransform` which utilizes `jax.experimental.io_callback` under the hood.
+
+```python
+import jax.numpy as jnp
+import numpy as np
+
+from cyreal import (
+  ArraySampleSource,
+  BatchTransform,
+  DataLoader,
+  HostCallbackTransform,
+  MNISTDataset,
+)
+
+def model(images):
+  return jnp.mean(images.astype(jnp.float32), axis=(1, 2, 3))
+
+def cross_entropy(logits, labels):
+  labels = labels.astype(jnp.float32)
+  return (logits - labels) ** 2
+
+def log_loss(batch, mask):
+  logits = model(batch["image"])
+  loss = jnp.mean(cross_entropy(logits, batch["label"]) * mask[:, None])
+  print("loss:", float(np.asarray(loss)))
+  return batch
+
+loader = DataLoader(
+  pipeline=[
+    ArraySampleSource(MNISTDataset(split="train").as_array_dict(), ordering="shuffle"),
+    BatchTransform(batch_size=128),
+    HostCallbackTransform(fn=log_loss),
+  ],
+)
+```

cyreal-0.1.1/cyreal/__init__.py

@@ -0,0 +1,38 @@
+"""Jittable dataset utilities for JAX."""
+from __future__ import annotations
+
+from .dataset_protocol import DatasetProtocol
+from .datasets import CIFAR10Dataset, CIFAR10DiskSource, MNISTDataset, MNISTDiskSource
+from .loader import (
+    DataLoader,
+    LoaderState,
+)
+from .sources import ArraySampleSource, DiskSampleSource, GymnaxSource, Source
+from .transforms import (
+    BatchTransform,
+    DevicePutTransform,
+    FlattenTransform,
+    HostCallbackTransform,
+    MapTransform,
+    NormalizeImageTransform,
+)
+
+__all__ = [
+    "DatasetProtocol",
+    "DataLoader",
+    "LoaderState",
+    "CIFAR10Dataset",
+    "CIFAR10DiskSource",
+    "MNISTDataset",
+    "MNISTDiskSource",
+    "ArraySampleSource",
+    "DiskSampleSource",
+    "GymnaxSource",
+    "Source",
+    "BatchTransform",
+    "DevicePutTransform",
+    "FlattenTransform",
+    "HostCallbackTransform",
+    "MapTransform",
+    "NormalizeImageTransform",
+]

cyreal-0.1.1/cyreal/dataset_protocol.py

@@ -0,0 +1,14 @@
+"""Simple dataset protocol used by the jittable dataloader."""
+from __future__ import annotations
+
+from typing import Any, Protocol
+
+
+class DatasetProtocol(Protocol):
+    """Minimal interface for indexable, length-known datasets."""
+
+    def __len__(self) -> int:
+        ...
+
+    def __getitem__(self, index: int) -> Any:
+        ...

cyreal-0.1.1/cyreal/datasets/__init__.py

@@ -0,0 +1,14 @@
+"""Dataset helpers bundled with cyreal."""
+from __future__ import annotations
+
+from .cifar10 import CIFAR10Dataset, CIFAR10DiskSource
+from .mnist import MNISTDataset, MNISTDiskSource, MNIST_URLS, _ensure_file
+
+__all__ = [
+    "CIFAR10Dataset",
+    "CIFAR10DiskSource",
+    "MNISTDataset",
+    "MNISTDiskSource",
+    "MNIST_URLS",
+    "_ensure_file",
+]

cyreal-0.1.1/cyreal/datasets/cifar10.py

@@ -0,0 +1,194 @@
+"""CIFAR-10 dataset that stays within NumPy/JAX dependencies."""
+from __future__ import annotations
+
+import pickle
+import tarfile
+import urllib.request
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Literal
+
+import jax
+import jax.numpy as jnp
+import numpy as np
+
+from ..dataset_protocol import DatasetProtocol
+from ..sources import DiskSampleSource
+
+CIFAR10_URL = "https://www.cs.toronto.edu/~kriz/cifar-10-python.tar.gz"
+
+
+def _to_host_jax_array(array: np.ndarray) -> jax.Array:
+    cpu_devices = jax.devices("cpu")
+    if cpu_devices:
+        with jax.default_device(cpu_devices[0]):
+            return jnp.asarray(array)
+    return jnp.asarray(array)
+
+
+def _resolve_cache_dir(cache_dir: str | Path | None) -> Path:
+    base_dir = Path(cache_dir) if cache_dir is not None else Path.home() / ".cache" / "cyreal_cifar10"
+    base_dir.mkdir(parents=True, exist_ok=True)
+    return base_dir
+
+
+def _download(url: str, path: Path) -> Path:
+    if path.exists():
+        return path
+    path.parent.mkdir(parents=True, exist_ok=True)
+    urllib.request.urlretrieve(url, path)
+    return path
+
+
+def _ensure_extracted(archive: Path, extract_root: Path) -> Path:
+    target = extract_root / "cifar-10-batches-py"
+    if target.exists():
+        return target
+    with tarfile.open(archive, "r:gz") as tar:
+        tar.extractall(path=extract_root)
+    return target
+
+
+def _batch_names(split: Literal["train", "test"]) -> list[str]:
+    if split == "train":
+        return [f"data_batch_{i}" for i in range(1, 6)]
+    if split == "test":
+        return ["test_batch"]
+    raise ValueError("split must be 'train' or 'test'.")
+
+
+def _load_split_numpy(split: Literal["train", "test"], extract_dir: Path) -> tuple[np.ndarray, np.ndarray]:
+    images = []
+    labels = []
+    for name in _batch_names(split):
+        batch_path = extract_dir / name
+        if not batch_path.exists():
+            raise FileNotFoundError(f"Missing CIFAR-10 batch '{name}'.")
+        with open(batch_path, "rb") as f:
+            batch = pickle.load(f, encoding="latin1")
+        data = batch["data"].reshape(-1, 3, 32, 32)
+        images.append(np.transpose(data, (0, 2, 3, 1)).astype(np.uint8))
+        labels.append(np.asarray(batch["labels"], dtype=np.int32))
+
+    images_np = np.concatenate(images, axis=0)
+    labels_np = np.concatenate(labels, axis=0)
+    return images_np, labels_np
+
+
+def _ensure_split_numpy_cache(
+    split: Literal["train", "test"],
+    base_dir: Path,
+    extract_dir: Path,
+) -> tuple[Path, Path]:
+    cache_root = base_dir / "disk_cache"
+    cache_root.mkdir(parents=True, exist_ok=True)
+    images_path = cache_root / f"{split}_images.npy"
+    labels_path = cache_root / f"{split}_labels.npy"
+    if not images_path.exists() or not labels_path.exists():
+        images_np, labels_np = _load_split_numpy(split, extract_dir)
+        np.save(images_path, images_np)
+        np.save(labels_path, labels_np)
+    return images_path, labels_path
+
+
+@dataclass
+class CIFAR10Dataset(DatasetProtocol):
+    """Download-free CIFAR-10 access ready for `ArraySampleSource`."""
+
+    split: Literal["train", "test"] = "train"
+    cache_dir: str | Path | None = None
+
+    def __post_init__(self) -> None:
+        base_dir = _resolve_cache_dir(self.cache_dir)
+        archive_path = base_dir / "cifar-10-python.tar.gz"
+        _download(CIFAR10_URL, archive_path)
+        extract_dir = _ensure_extracted(archive_path, base_dir)
+
+        images_np, labels_np = _load_split_numpy(self.split, extract_dir)
+        self._images = _to_host_jax_array(images_np)
+        self._labels = _to_host_jax_array(labels_np)
+
+    def __len__(self) -> int:
+        return int(self._images.shape[0])
+
+    def __getitem__(self, index: int):
+        return {
+            "image": self._images[index],
+            "label": self._labels[index],
+        }
+
+    def as_array_dict(self) -> dict[str, jax.Array]:
+        """Expose the full dataset as a PyTree of JAX arrays."""
+
+        return {
+            "image": self._images,
+            "label": self._labels,
+        }
+
+    @classmethod
+    def make_disk_source(
+        cls,
+        *,
+        split: Literal["train", "test"] = "train",
+        cache_dir: str | Path | None = None,
+        ordering: Literal["sequential", "shuffle"] = "shuffle",
+        prefetch_size: int = 64,
+    ) -> DiskSampleSource:
+        base_dir = _resolve_cache_dir(cache_dir)
+        archive_path = base_dir / "cifar-10-python.tar.gz"
+        _download(CIFAR10_URL, archive_path)
+        extract_dir = _ensure_extracted(archive_path, base_dir)
+
+        images_path, labels_path = _ensure_split_numpy_cache(split, base_dir, extract_dir)
+        images_memmap = np.load(images_path, mmap_mode="r")
+        labels_memmap = np.load(labels_path, mmap_mode="r")
+
+        if images_memmap.shape[0] != labels_memmap.shape[0]:
+            raise ValueError("CIFAR-10 image and label counts do not match.")
+
+        def _read_sample(index: int | np.ndarray) -> dict[str, np.ndarray]:
+            idx = int(np.asarray(index))
+            image = np.asarray(images_memmap[idx], dtype=np.uint8)
+            label = np.asarray(labels_memmap[idx], dtype=np.int32)
+            return {"image": image, "label": label}
+
+        sample_spec = {
+            "image": jax.ShapeDtypeStruct(shape=tuple(images_memmap.shape[1:]), dtype=jnp.uint8),
+            "label": jax.ShapeDtypeStruct(shape=(), dtype=jnp.int32),
+        }
+
+        return DiskSampleSource(
+            length=int(labels_memmap.shape[0]),
+            sample_fn=_read_sample,
+            sample_spec=sample_spec,
+            ordering=ordering,
+            prefetch_size=prefetch_size,
+        )
+
+
+@dataclass
+class CIFAR10DiskSource:
+    """Stream CIFAR-10 samples directly from cached numpy arrays on disk."""
+
+    split: Literal["train", "test"] = "train"
+    cache_dir: str | Path | None = None
+    ordering: Literal["sequential", "shuffle"] = "shuffle"
+    prefetch_size: int = 64
+
+    def __post_init__(self) -> None:
+        self._disk_source = CIFAR10Dataset.make_disk_source(
+            split=self.split,
+            cache_dir=self.cache_dir,
+            ordering=self.ordering,
+            prefetch_size=self.prefetch_size,
+        )
+        self.steps_per_epoch = self._disk_source.steps_per_epoch
+
+    def element_spec(self):
+        return self._disk_source.element_spec()
+
+    def init_state(self, key=None):
+        return self._disk_source.init_state(key)
+
+    def next(self, state):
+        return self._disk_source.next(state)