jax example

.darglint

@@ -0,0 +1,3 @@
+[darglint]
+
+docstring_style=google

.github/workflows/publish.yml

@@ -0,0 +1,47 @@
+name: Publish Python Package
+
+on:
+  release:
+    types: [created]
+  workflow_dispatch:
+
+permissions:
+  contents: read
+  id-token: write
+
+concurrency:
+  group: "publish"
+  cancel-in-progress: true
+
+jobs:
+  publish:
+    timeout-minutes: 10
+    name: Build and publish
+
+    # We don't need to run on all platforms since this package is
+    # platform-agnostic. The output wheel is something like
+    # "monotonic_attention-<version>-py3-none-any.whl".
+    runs-on: ubuntu-latest
+
+    steps:
+      - name: Checkout code
+        uses: actions/checkout@v3
+
+      - name: Set up Python
+        uses: actions/setup-python@v4
+        with:
+          python-version: "3.10"
+
+      - name: Install dependencies
+        run: |
+          python -m pip install --upgrade pip
+          pip install build wheel
+
+      - name: Build package
+        run: python -m build --sdist --wheel --outdir dist/ .
+
+      - name: Publish package
+        uses: pypa/gh-action-pypi-publish@release/v1
+        # with:
+        #   user: __token__
+        #   password: ${{ secrets.PYPI_API_TOKEN }}

.github/workflows/test.yml

@@ -0,0 +1,65 @@
+name: Python Checks
+
+on:
+  push:
+    branches:
+      - master
+  pull_request:
+    branches:
+      - master
+    types:
+      - opened
+      - reopened
+      - synchronize
+      - ready_for_review
+
+concurrency:
+  group: tests-${{ github.head_ref || github.run_id }}
+  cancel-in-progress: true
+
+jobs:
+  run-base-tests:
+    timeout-minutes: 10
+    runs-on: ubuntu-latest
+    steps:
+      - name: Check out repository
+        uses: actions/checkout@v3
+
+      - name: Set up Python
+        uses: actions/setup-python@v4
+        with:
+          python-version: "3.11"
+
+      - name: Restore cache
+        id: restore-cache
+        uses: actions/cache/restore@v3
+        with:
+          path: |
+            ${{ env.pythonLocation }}
+            .mypy_cache/
+          key: python-requirements-${{ env.pythonLocation }}-${{ github.event.pull_request.base.sha || github.sha }}
+          restore-keys: |
+            python-requirements-${{ env.pythonLocation }}
+            python-requirements-
+
+      - name: Install package
+        run: |
+          pip install --upgrade --upgrade-strategy eager '.[dev]' 'jax[cpu]'
+
+      - name: Run static checks
+        run: |
+          mkdir -p .mypy_cache
+          make static-checks
+
+      - name: Run unit tests
+        run: |
+          make test
+
+      - name: Save cache
+        uses: actions/cache/save@v3
+        if: github.ref == 'refs/heads/master'
+        with:
+          path: |
+            ${{ env.pythonLocation }}
+            .mypy_cache/
+          key: ${{ steps.restore-cache.outputs.cache-primary-key }}

.gitignore

@@ -0,0 +1,21 @@
+# .gitignore
+
+# Python
+*.py[oc]
+__pycache__/
+*.egg-info
+.eggs/
+.mypy_cache/*
+.pyre/
+.pytest_cache/
+.ruff_cache/
+.dmypy.json
+
+# Databases
+*.db
+
+# Build artifacts
+build/
+dist/
+*.so
+out*/

LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2023 Benjamin Bolte
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

MANIFEST.in

@@ -0,0 +1 @@
+recursive-include xax/ *.py *.txt py.typed MANIFEST.in

Makefile

@@ -0,0 +1,77 @@
+# Makefile
+
+define HELP_MESSAGE
+                    xax
+
+# Installing
+
+1. Create a new Conda environment: `conda create --name xax python=3.11`
+2. Activate the environment: `conda activate xax`
+3. Install the package: `make install-dev`
+
+# Running Tests
+
+1. Run autoformatting: `make format`
+2. Run static checks: `make static-checks`
+3. Run unit tests: `make test`
+
+endef
+export HELP_MESSAGE
+
+all:
+	@echo "$$HELP_MESSAGE"
+.PHONY: all
+
+# ------------------------ #
+#          Build           #
+# ------------------------ #
+
+install:
+	@pip install --verbose -e .
+.PHONY: install
+
+install-dev:
+	@pip install --verbose -e '.[dev]'
+.PHONY: install
+
+build-ext:
+	@python setup.py build_ext --inplace
+.PHONY: build-ext
+
+clean:
+	rm -rf build dist *.so **/*.so **/*.pyi **/*.pyc **/*.pyd **/*.pyo **/__pycache__ *.egg-info .eggs/ .ruff_cache/
+.PHONY: clean
+
+# ------------------------ #
+#       Static Checks      #
+# ------------------------ #
+
+py-files := $(shell find . -name '*.py')
+
+format:
+	@black $(py-files)
+	@ruff --fix $(py-files)
+.PHONY: format
+
+format-cpp:
+	@clang-format -i $(shell find . -name '*.cpp' -o -name '*.h')
+	@cmake-format -i $(shell find . -name 'CMakeLists.txt' -o -name '*.cmake')
+.PHONY: format-cpp
+
+static-checks:
+	@black --diff --check $(py-files)
+	@ruff $(py-files)
+	@mypy --install-types --non-interactive $(py-files)
+.PHONY: lint
+
+mypy-daemon:
+	@dmypy run -- $(py-files)
+.PHONY: mypy-daemon
+
+# ------------------------ #
+#        Unit tests        #
+# ------------------------ #
+
+test:
+	python -m pytest
+.PHONY: test

README.md

@@ -1,2 +1,3 @@
 # xax
-JAX library for fast experimentation
+
+JAX library for fast experimentation.

examples/mnist.py

@@ -0,0 +1,148 @@
+"""MNIST example in Jax."""
+
+import array
+import gzip
+import itertools
+import os
+import struct
+import time
+import urllib.request
+from typing import Iterator
+
+import jax.numpy as jnp
+import numpy as np
+import numpy.random as npr
+from jax import grad, jit, random
+from jax.example_libraries import optimizers, stax
+from jax.example_libraries.optimizers import OptimizerState
+from jax.example_libraries.stax import Dense, LogSoftmax, Relu
+from jaxtyping import ArrayLike
+
+_DATA = "/tmp/jax_example_data/"
+
+
+def _download(url: str, filename: str) -> None:
+    if not os.path.exists(_DATA):
+        os.makedirs(_DATA)
+    out_file = os.path.join(_DATA, filename)
+    if not os.path.isfile(out_file):
+        urllib.request.urlretrieve(url, out_file)
+        print(f"downloaded {url} to {_DATA}")
+
+
+def _partial_flatten(x: np.ndarray) -> np.ndarray:
+    return np.reshape(x, (x.shape[0], -1))
+
+
+def _one_hot(x: np.ndarray, k: int, dtype: type = np.float32) -> np.ndarray:
+    return np.array(x[:, None] == np.arange(k), dtype)
+
+
+def mnist_raw() -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
+    base_url = "https://storage.googleapis.com/cvdf-datasets/mnist/"
+
+    def parse_labels(filename: str) -> np.ndarray:
+        with gzip.open(filename, "rb") as fh:
+            _ = struct.unpack(">II", fh.read(8))
+            return np.array(array.array("B", fh.read()), dtype=np.uint8)
+
+    def parse_images(filename: str) -> np.ndarray:
+        with gzip.open(filename, "rb") as fh:
+            _, num_data, rows, cols = struct.unpack(">IIII", fh.read(16))
+            return np.array(array.array("B", fh.read()), dtype=np.uint8).reshape(num_data, rows, cols)
+
+    for filename in [
+        "train-images-idx3-ubyte.gz",
+        "train-labels-idx1-ubyte.gz",
+        "t10k-images-idx3-ubyte.gz",
+        "t10k-labels-idx1-ubyte.gz",
+    ]:
+        _download(base_url + filename, filename)
+
+    train_images = parse_images(os.path.join(_DATA, "train-images-idx3-ubyte.gz"))
+    train_labels = parse_labels(os.path.join(_DATA, "train-labels-idx1-ubyte.gz"))
+    test_images = parse_images(os.path.join(_DATA, "t10k-images-idx3-ubyte.gz"))
+    test_labels = parse_labels(os.path.join(_DATA, "t10k-labels-idx1-ubyte.gz"))
+
+    return train_images, train_labels, test_images, test_labels
+
+
+def mnist(permute_train: bool = False) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
+    train_images, train_labels, test_images, test_labels = mnist_raw()
+
+    train_images = _partial_flatten(train_images) / np.float32(255.0)
+    test_images = _partial_flatten(test_images) / np.float32(255.0)
+    train_labels = _one_hot(train_labels, 10)
+    test_labels = _one_hot(test_labels, 10)
+
+    if permute_train:
+        perm = np.random.RandomState(0).permutation(train_images.shape[0])
+        train_images = train_images[perm]
+        train_labels = train_labels[perm]
+
+    return train_images, train_labels, test_images, test_labels
+
+
+def loss(params: tuple[ArrayLike, ArrayLike], batch: tuple[ArrayLike, ArrayLike]) -> ArrayLike:
+    inputs, targets = batch
+    preds = predict(params, inputs)
+    return -jnp.mean(jnp.sum(preds * targets, axis=1))
+
+
+def accuracy(params: tuple[ArrayLike, ArrayLike], batch: tuple[ArrayLike, ArrayLike]) -> ArrayLike:
+    inputs, targets = batch
+    target_class = jnp.argmax(targets, axis=1)
+    predicted_class = jnp.argmax(predict(params, inputs), axis=1)
+    return jnp.mean(predicted_class == target_class)
+
+
+init_random_params, predict = stax.serial(Dense(1024), Relu, Dense(1024), Relu, Dense(10), LogSoftmax)
+
+if __name__ == "__main__":
+    # python -m examples.mnist
+    rng = random.PRNGKey(0)
+
+    step_size = 0.001
+    num_epochs = 10
+    batch_size = 128
+    momentum_mass = 0.9
+
+    train_images, train_labels, test_images, test_labels = mnist()
+    num_train = train_images.shape[0]
+    num_complete_batches, leftover = divmod(num_train, batch_size)
+    num_batches = num_complete_batches + bool(leftover)
+
+    def data_stream() -> Iterator[tuple[np.ndarray, np.ndarray]]:
+        rng = npr.RandomState(0)
+        while True:
+            perm = rng.permutation(num_train)
+            for i in range(num_batches):
+                batch_idx = perm[i * batch_size : (i + 1) * batch_size]
+                yield train_images[batch_idx], train_labels[batch_idx]
+
+    batches = data_stream()
+
+    opt_init, opt_update, get_params = optimizers.momentum(step_size, mass=momentum_mass)
+
+    @jit
+    def update(i: int, opt_state: OptimizerState, batch: tuple[ArrayLike, ArrayLike]) -> OptimizerState:
+        params = get_params(opt_state)
+        return opt_update(i, grad(loss)(params, batch), opt_state)
+
+    _, init_params = init_random_params(rng, (-1, 28 * 28))
+    opt_state = opt_init(init_params)
+    itercount = itertools.count()
+
+    print("\nStarting training...")
+    for epoch in range(num_epochs):
+        start_time = time.time()
+        for _ in range(num_batches):
+            opt_state = update(next(itercount), opt_state, next(batches))
+        epoch_time = time.time() - start_time
+
+        params = get_params(opt_state)
+        train_acc = accuracy(params, (train_images, train_labels))
+        test_acc = accuracy(params, (test_images, test_labels))
+        print(f"Epoch {epoch} in {epoch_time:0.2f} sec")
+        print(f"Training set accuracy {train_acc}")
+        print(f"Test set accuracy {test_acc}")