Switch to OxCaml

.github/workflows/build.yml

@@ -0,0 +1,77 @@
+name: Build and Test
+
+on:
+  push:
+    branches: [main, master]
+  pull_request:
+    branches: [main, master]
+
+jobs:
+  build:
+    # xla binds a prebuilt XLA extension blob; only the linux x86_64 archive is
+    # downloaded below, so this job is Linux-only (macOS would need the darwin archive).
+    runs-on: ubuntu-latest
+
+    env:
+      TZ: UTC
+
+    steps:
+      # The OxCaml + Jane Street tree is large and overflows the ~14 GB free on a stock
+      # ubuntu-latest runner; strip preinstalled tooling we don't need to recover ~30 GB.
+      # tool-cache:false keeps the hosted tool cache (used by setup-ocaml).
+      - name: Free up disk space
+        uses: jlumbroso/free-disk-space@v1.3.1
+        with:
+          tool-cache: false
+          android: true
+          dotnet: true
+          haskell: true
+          large-packages: true
+          docker-images: true
+          swap-storage: true
+
+      - name: Checkout code
+        uses: actions/checkout@v6
+
+      - name: Set-up OxCaml
+        uses: ocaml/setup-ocaml@v3
+        with:
+          ocaml-compiler: ocaml-variants.5.2.0+ox
+          dune-cache: true
+          opam-pin: false
+          opam-repositories: |
+            oxcaml: "git+https://github.com/oxcaml/opam-repository.git"
+            default: "git+https://github.com/ocaml/opam-repository.git"
+
+      # xla's C++ stubs need the XLA extension headers at *opam install* time (not just at
+      # dune build time), so fetch the blob and export XLA_EXTENSION_DIR before installing.
+      - name: Download XLA extension
+        run: |
+          wget -q https://github.com/elixir-nx/xla/releases/download/v0.4.4/xla_extension-x86_64-linux-gnu-cpu.tar.gz
+          tar -xzf xla_extension-x86_64-linux-gnu-cpu.tar.gz
+          echo "XLA_EXTENSION_DIR=$PWD/xla_extension" >> "$GITHUB_ENV"
+
+      # The compiled dependencies live in the local _opam switch, not ~/.opam; caching only
+      # ~/.opam would leave every run recompiling the whole tree. Both paths are cached,
+      # keyed on the lockfile, and refreshed automatically in the post-job step.
+      - name: Cache opam switch
+        uses: actions/cache@v5
+        with:
+          path: |
+            ~/.opam
+            _opam
+          key: opam-${{ hashFiles('fox.opam.locked') }}
+          restore-keys: |
+            opam-
+
+      - name: Install opam dependencies
+        run: opam install . --deps-only --with-test --locked
+
+      - name: Build
+        run: opam exec -- dune build @default
+
+      - name: Run tests
+        run: opam exec -- dune runtest
+
+      - name: Check formatting
+        run: opam exec -- dune build @fmt

.gitignore

@@ -1 +1,5 @@
 _build
+_opam
+*.install
+xla_extension
+/mnist

.ocamlformat

@@ -1,3 +1 @@
 profile = janestreet
-version = 0.27.0
-ocaml-version = 5.3

README.md

@@ -18,8 +18,48 @@ many features.
 
 ## Building
 
+Fox targets [OxCaml](https://oxcaml.org/): it uses OxCaml's effect-handler
+syntax and the modal Jane Street libraries, and builds against a local opam
+switch on `ocaml-variants.5.2.0+ox` with a committed lockfile
+(`fox.opam.locked`). It depends on a pinned
+[`xla`](https://github.com/mt-caret/ocaml-xla), which binds a prebuilt XLA
+extension blob.
+
+### One-time setup
+
+Download the XLA extension (elixir-nx/xla v0.4.4) and point `XLA_EXTENSION_DIR`
+at it. The C++ stubs need its headers at build time; the runtime `rpath` is
+baked in, so no `LD_LIBRARY_PATH` is needed once built:
+
+```bash
+wget https://github.com/elixir-nx/xla/releases/download/v0.4.4/xla_extension-x86_64-linux-gnu-cpu.tar.gz
+tar -xzf xla_extension-x86_64-linux-gnu-cpu.tar.gz   # -> ./xla_extension
+export XLA_EXTENSION_DIR=$PWD/xla_extension
+```
+
+(On platforms other than linux x86_64, download the matching archive instead.)
+
+Create the switch from the lockfile:
+
+```bash
+opam switch create . 5.2.0+ox \
+  --repos ox=git+https://github.com/oxcaml/opam-repository.git,default \
+  --locked
+```
+
+### Build and test
+
+```bash
+XLA_EXTENSION_DIR=$PWD/xla_extension dune build @default @runtest
+```
+
+### Regenerating the lockfile
+
+After changing dependencies, refresh `fox.opam.locked`:
+
 ```bash
-DYLD_LIBRARY_PATH=/absolute-path-to-ocaml-xla-extension-lib dune build @default @runtest -w
+XLA_EXTENSION_DIR=$PWD/xla_extension opam install . --deps-only --with-test
+opam lock .
 ```
 
 ## TODO

dune

@@ -0,0 +1,13 @@
+; XLA's C++ runtime logs to stdout/stderr the first time it initialises a CPU
+; client, which pollutes inline-test output non-deterministically. Silence
+; everything below ERROR so expect-test output is stable.
+;
+; oneDNN reorders floating-point reductions per host CPU, which perturbs the low
+; bits of XLA results and makes exact expect tests machine-dependent. Disable it
+; so numeric outputs are reproducible across machines (e.g. local vs CI).
+
+(env
+ (_
+  (env-vars
+   (TF_CPP_MIN_LOG_LEVEL 2)
+   (TF_ENABLE_ONEDNN_OPTS 0))))

dune-project

@@ -1,3 +1,39 @@
-(lang dune 3.17)
+(lang dune 3.20)
 
 (name fox)
+
+(generate_opam_files true)
+
+(source
+ (github mt-caret/fox))
+
+(authors "mt-caret")
+
+(maintainers "mt-caret")
+
+(package
+ (name fox)
+ (synopsis "Autodiff and XLA JIT compilation for OCaml, inspired by JAX")
+ (description
+  "Fox is an OCaml library for automatic differentiation and just-in-time XLA compilation of tensor computations, inspired by JAX.")
+ (allow_empty)
+ (depends
+  core
+  core_kernel
+  core_unix
+  base_quickcheck
+  expect_test_helpers_core
+  ppx_jane
+  ppx_typed_fields
+  splittable_random
+  xla
+  (ocamlformat :with-dev-setup)
+  (odoc
+   (or :with-doc :with-dev-setup))
+  (utop :with-dev-setup)))
+
+(pin
+ (package
+  (name xla))
+ (url
+  "git+https://github.com/mt-caret/ocaml-xla.git#6dffd132a9e1b13b3b8c17f1abae110c40ea6f3c"))