diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml new file mode 100644 index 0000000..9d89e00 --- /dev/null +++ b/.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 diff --git a/.gitignore b/.gitignore index e35d885..6b8c4df 100644 --- a/.gitignore +++ b/.gitignore @@ -1 +1,5 @@ _build +_opam +*.install +xla_extension +/mnist diff --git a/.ocamlformat b/.ocamlformat index 6a9437d..0619e80 100644 --- a/.ocamlformat +++ b/.ocamlformat @@ -1,3 +1 @@ profile = janestreet -version = 0.27.0 -ocaml-version = 5.3 diff --git a/README.md b/README.md index dc329c4..3a0d987 100644 --- a/README.md +++ b/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 diff --git a/dune b/dune new file mode 100644 index 0000000..dd45564 --- /dev/null +++ b/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)))) diff --git a/dune-project b/dune-project index 277df53..8064d0e 100644 --- a/dune-project +++ b/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")) diff --git a/example/mnist.ml b/example/mnist.ml index 3679542..69d6129 100644 --- a/example/mnist.ml +++ b/example/mnist.ml @@ -23,26 +23,26 @@ module Dataset = struct in Tensor.init Float - ~dims:[| Bigstring.length contents |] - ~f:(fun index -> Bigstring.get contents index.(0) |> Char.to_int |> Int.to_float) + ~dims:[: Bigstring.length contents :] + ~f:(fun index -> Bigstring.get contents index.:(0) |> Char.to_int |> Int.to_float) ;; let load ~x ~y = - { x = load ~header_length:16 x |> Tensor.reshape ~dims:[| -1; 28; 28 |] + { x = load ~header_length:16 x |> Tensor.reshape ~dims:[: -1; 28; 28 :] ; y = load ~header_length:8 y } ;; let print { x; y } ~i = - let n = Tensor.get_exn Float y [| i |] |> Int.of_float in + let n = Tensor.get_exn Float y [: i :] |> Int.of_float in print_endline [%string "Label: %{n#Int}"]; - let x = Tensor.left_slice x ~indices:[| i |] in + let x = Tensor.left_slice x ~indices:[: i :] in let image = List.range 0 28 |> List.map ~f:(fun row -> List.range 0 28 |> List.map ~f:(fun col -> - let pixel = Tensor.get_exn Float x [| row; col |] in + let pixel = Tensor.get_exn Float x [: row; col :] in let pixel = 255 - Int.of_float (pixel /. 255. *. 23.) in [%string "\027[48;5;%{pixel#Int}m "]) |> String.concat) @@ -53,40 +53,37 @@ module Dataset = struct end module Model = struct - module T = struct - type t = - { h1 : Value.t - ; b1 : Value.t - ; h2 : Value.t - ; b2 : Value.t - ; h3 : Value.t - ; b3 : Value.t - } - [@@deriving typed_fields, sexp_of] - - let field_treeable (type a) (field : a Typed_field.t) - : (a -> Value_tree.t) * (module Treeable.S with type t = a) - = - match field with - | H1 -> Value.tree_of_t, (module Value) - | B1 -> Value.tree_of_t, (module Value) - | H2 -> Value.tree_of_t, (module Value) - | B2 -> Value.tree_of_t, (module Value) - | H3 -> Value.tree_of_t, (module Value) - | B3 -> Value.tree_of_t, (module Value) - ;; - end - - include T - include Treeable.Of_typed_fields (T) + type t = + { h1 : Value.t + ; b1 : Value.t + ; h2 : Value.t + ; b2 : Value.t + ; h3 : Value.t + ; b3 : Value.t + } + [@@deriving typed_fields, sexp_of] + + let field_treeable (type a) (local_ (field : a Typed_field.t)) + : (a -> Value_tree.t) * (module Treeable.S with type t = a) + = + match field with + | H1 -> Value.tree_of_t, (module Value) + | B1 -> Value.tree_of_t, (module Value) + | H2 -> Value.tree_of_t, (module Value) + | B2 -> Value.tree_of_t, (module Value) + | H3 -> Value.tree_of_t, (module Value) + | B3 -> Value.tree_of_t, (module Value) + ;; + + include functor Treeable.Of_typed_fields let create ~rng = - { h1 = Tensor.normal ~dims:[| 784; 128 |] ~rng () |> Value.of_tensor - ; b1 = Tensor.zeros ~dims:[| 128 |] |> Value.of_tensor - ; h2 = Tensor.normal ~dims:[| 128; 64 |] ~rng () |> Value.of_tensor - ; b2 = Tensor.zeros ~dims:[| 64 |] |> Value.of_tensor - ; h3 = Tensor.normal ~dims:[| 64; 10 |] ~rng () |> Value.of_tensor - ; b3 = Tensor.zeros ~dims:[| 10 |] |> Value.of_tensor + { h1 = Tensor.normal ~dims:[: 784; 128 :] ~rng () |> Value.of_tensor + ; b1 = Tensor.zeros ~dims:[: 128 :] |> Value.of_tensor + ; h2 = Tensor.normal ~dims:[: 128; 64 :] ~rng () |> Value.of_tensor + ; b2 = Tensor.zeros ~dims:[: 64 :] |> Value.of_tensor + ; h3 = Tensor.normal ~dims:[: 64; 10 :] ~rng () |> Value.of_tensor + ; b3 = Tensor.zeros ~dims:[: 10 :] |> Value.of_tensor } ;; @@ -108,18 +105,18 @@ module Model = struct let linear ~h ~b ~bs x = let open Value.O in - let b = Value.broadcast b ~dims:(Array.append [| bs |] (Value.dims b)) in + let b = Value.broadcast b ~dims:(Iarray.append [: bs :] (Value.dims b)) in Value.matmul x h + b ;; let run { h1; b1; h2; b2; h3; b3 } x = match Value.dims x with - | [| bs; 784 |] -> + | [: bs; 784 :] -> let x (* bs x 128 *) = linear ~h:h1 ~b:b1 ~bs x |> Value.sigmoid in let x (* bs x 64 *) = linear ~h:h2 ~b:b2 ~bs x |> Value.sigmoid in let x (* bs x 10 *) = linear ~h:h3 ~b:b3 ~bs x |> Value.softmax ~dim:1 in x - | _ -> raise_s [%message "Invalid input dimensions" ~dims:(Value.dims x : int array)] + | _ -> raise_s [%message "Invalid input dimensions" ~dims:(Value.dims x : int iarray)] ;; let cross_entropy_loss t ~x ~y = @@ -152,60 +149,62 @@ let command = Dataset.print train ~i:0; let rng = Splittable_random.of_int seed in let model = ref (Model.create ~rng) in - let print_dataset_loss () = + let batch_values ~images ~labels = let x = - Tensor.reshape train.x ~dims:[| -1; 28 * 28 |] + Tensor.reshape images ~dims:[: -1; 28 * 28 :] |> Tensor.to_typed_exn Float |> Tensor.Typed.map Float ~f:(fun x -> x /. 255.) |> Value.of_typed_tensor in let y = - let labels = train.y in Tensor.Typed.init Float - ~dims:[| Tensor.length labels; 10 |] + ~dims:[: Tensor.length labels; 10 :] ~f:(fun index -> - let label = Tensor.get_exn Float labels [| index.(0) |] |> Float.to_int in - if label = index.(1) then 1. else 0.) + let label = Tensor.get_exn Float labels [: index.:(0) :] |> Float.to_int in + if label = index.:(1) then 1. else 0.) |> Value.of_typed_tensor in - let loss = - Fox_jit.jit - (module Model) - (module Value) - ~f:(fun model -> Model.cross_entropy_loss model ~x ~y) - !model - in + x, y + in + (* The [jit] closures are created once, here, rather than rebuilt per iteration. Each + takes its batch [(x, y)] as inputs, so the program is compiled on the first call + and the executable is reused for every later call of the same structure (all + batches share a shape). *) + let dataset_loss = + Staged.unstage + (Fox_jit.jit + (module Treeable.Tuple2 (Model) (Treeable.Tuple2 (Value) (Value))) + (module Value) + ~f:(fun (model, (x, y)) -> Model.cross_entropy_loss model ~x ~y) + ()) + in + let eval_x, eval_y = batch_values ~images:train.x ~labels:train.y in + let print_dataset_loss () = + let loss = dataset_loss (!model, (eval_x, eval_y)) in print_s [%message "test dataset loss" (loss : Value.t)] in + let train_step = + Staged.unstage + (Fox_jit.jit + (module Treeable.Tuple2 (Model) (Treeable.Tuple2 (Value) (Value))) + (module Treeable.Tuple2 (Value) (Model)) + ~f:(fun (model, (x, y)) -> + grad_and_value + (module Model) + ~f:(fun model -> Model.cross_entropy_loss model ~x ~y) + ~x:model) + ()) + in for i = 0 to (Dataset.length train / batch_size) - 1 do - let x = - Tensor.sub_left train.x ~pos:(i * batch_size) ~len:batch_size - |> Tensor.reshape ~dims:[| -1; 28 * 28 |] - |> Tensor.to_typed_exn Float - |> Tensor.Typed.map Float ~f:(fun x -> x /. 255.) - |> Value.of_typed_tensor - in - let y = - let labels = Tensor.sub_left train.y ~pos:(i * batch_size) ~len:batch_size in - Tensor.Typed.init Float ~dims:[| batch_size; 10 |] ~f:(fun index -> - let label = Tensor.get_exn Float labels [| index.(0) |] |> Float.to_int in - if label = index.(1) then 1. else 0.) - |> Value.of_typed_tensor - in - let loss, grad = - Fox_jit.jit - (module Model) - (module Treeable.Tuple2 (Value) (Model)) - ~f:(fun model -> - grad_and_value - (module Model) - ~f:(fun model -> Model.cross_entropy_loss model ~x ~y) - ~x:model) - !model + let x, y = + batch_values + ~images:(Tensor.sub_left train.x ~pos:(i * batch_size) ~len:batch_size) + ~labels:(Tensor.sub_left train.y ~pos:(i * batch_size) ~len:batch_size) in + let loss, grad = train_step (!model, (x, y)) in let average_grad_l2_norm = - Eval.handle ~f:(fun () -> + eval ~f:(fun () -> let grad_norms = Model.map grad ~f:(fun x -> Value.mean Value.O.(x * x)) |> Model.to_list in @@ -215,7 +214,7 @@ let command = if i mod 100 = 0 then print_dataset_loss (); print_s [%message (loss : Value.t) (average_grad_l2_norm : Value.t)]; model - := Eval.handle ~f:(fun () -> + := eval ~f:(fun () -> Model.map2 !model grad ~f:(fun a b -> Value.O.(a - Value.scale b learning_rate))) done; diff --git a/fox-oxcaml-handoff.md b/fox-oxcaml-handoff.md new file mode 100644 index 0000000..0356759 --- /dev/null +++ b/fox-oxcaml-handoff.md @@ -0,0 +1,239 @@ +# Handoff: building `fox` with OxCaml + a reproducible opam lockfile + +## The goal + +Get **`fox`** (https://github.com/mt-caret/fox — a JAX-like OCaml library for autodiff + +XLA JIT) building with **OxCaml** and a **reproducible opam lock file**, mirroring the +setup of **`mt-caret/orc`** (the reference to copy). + +References: +- OxCaml: https://oxcaml.org/ +- opam lockfiles: https://ocamlpro.com/blog/2026_01_08_opam_104_sharing_your_code/#lockdependencies +- Pattern to mirror: https://github.com/mt-caret/orc + +The orc pattern = a **local `_opam` switch** on `ocaml-variants.5.2.0+ox`, `dune-project` +with `(generate_opam_files true)` + a `(package …)` stanza + `(pin …)` stanzas, a +`.opam.template` (for `pin-depends` + `x-maintenance-intent`), a committed +`.opam.locked`, and a GitHub Actions CI workflow. Create switches with +`opam switch create . 5.2.0+ox --repos ox=git+https://github.com/oxcaml/opam-repository.git,default --locked`. + +## My (the user's) preferences + +- **Jane Street / OxCaml OCaml style** per the global `~/.claude/CLAUDE.md`. Most relevant: + - Always `dune build` **and** `dune runtest` before claiming done; `dune fmt` for + formatting but **inspect the diff** before promoting. + - `open Core`; prefer Core's `List`/`Option`/`String`/etc. + - `.mli` for every non-test `.ml` except the library-name wrapper module. + - Don't introduce single-use abstractions (2+ uses is fine — e.g. the shared effect + handler helper below). + - Comments must state non-obvious things and stand alone for a fresh reader — **no** + "ported from X" / "now uses Y" history comments. + - `[%sexp_of: T]`/`[%of_sexp: T]` extension forms, ppx_let prefix forms, etc. +- **Git: leave changes uncommitted on `main`.** Don't commit or push unless I ask. +- I already approved doing the **full** job end-to-end (scaffold → switch → deps → build/ + test → lockfile → CI/README). The blocker below (modes porting) was hit mid-way. + +## What this task actually is (read this first) + +It is **not** just build config. `fox`'s source was written for **vanilla OCaml 5.3 + +non-modal Jane Street libs**, but **OxCaml is 5.2-based and the JS `v0.18~preview` libs are +modal** (local/global modes). The `ox` opam repo only ships `ocaml-variants.5.2.0+ox` +(there is **no 5.3+ox**). So you must **port `fox`'s source** to compile on 5.2.0+ox + the +modal libs. Expect to touch the library's core. Raise this scope with me if it balloons. + +## Config scaffolding (these all worked — recreate them) + +**`dune-project`:** +```scheme +(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 "…") + (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#main")) +``` +Dep mapping from the `dune` files: `core_kernel.nonempty_list`→`core_kernel`; +`core_unix.{bigstring_unix,command_unix}`→`core_unix`; `ppx_typed_fields.typed_fields_lib`+ +`base_quickcheck.ppx_quickcheck` are covered by `ppx_typed_fields`/`base_quickcheck`; +`ppx_jane` covers ppx_expect/inline_test. `xla` is the only non-JS dep. + +**`fox.opam.template`:** +``` +pin-depends: [ + ["xla.dev" "git+https://github.com/mt-caret/ocaml-xla.git#main"] +] +x-maintenance-intent: ["(latest)"] +``` +Then `dune build ./fox.opam` regenerates `fox.opam` (the `(pin …)` stanza is accepted by +dune 3.20; `pin-depends` comes through from the template). + +**`.ocamlformat`:** reduce to just `profile = janestreet`. The old file pinned +`version = 0.27.0` / `ocaml-version = 5.3`, which **conflicts** with OxCaml's +`ocamlformat 0.26.2+ox1` and breaks `dune build @fmt`. (The +ox ocamlformat DOES parse +OxCaml syntax. After porting, run `dune fmt` and inspect/`dune promote` the reflow.) + +**`.gitignore`:** add `_opam`, `*.install`, `xla_extension`. + +## The `xla` dependency + +- `xla` is **unpublished** (not in `default` or `ox` repos). Pin + `git+https://github.com/mt-caret/ocaml-xla.git#main` (mt-caret's fork of + LaurentMazare/ocaml-xla; default branch `main`). +- It binds a **prebuilt XLA extension blob**. Download elixir-nx/xla **v0.4.4**: + ```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/{include,lib/libxla_extension.so (~305MB)} + export XLA_EXTENSION_DIR=$PWD/xla_extension + ``` + (Other platforms: pick the matching archive.) The configurator + (`src/config/discover.ml`) resolves it via `XLA_EXTENSION_DIR` > `$DUNE_SOURCEROOT/ + xla_extension` > `$OPAM_SWITCH_PREFIX/lib/libxla`, and bakes `-Wl,-rpath,…/lib`, so **no + `LD_LIBRARY_PATH` is needed at runtime once built**. +- **`XLA_EXTENSION_DIR` must be set during `opam install` too** (the `xla_stubs.cpp` C++ + stubs need the headers at build time), not only for `dune build`. + +## ⚠️ Biggest pitfall: `mt-caret/ocaml-xla` has broken opam metadata + +Its `xla.opam` / `dune-project (depends …)` only lists +`base stdio camlzip ctypes ctypes-foreign ocaml dune`, but its dune files actually use: +- `src/wrapper/dune` libraries: **`int_repr`, `yojson`** (and `bigarray`); preprocess pps: + **`ppx_expect ppx_sexp_conv ppx_sexp_message ppx_compare`**. +- `src/config/dune`: `dune.configurator` (→ **`dune-configurator`**). + +Observed failures on a fresh install: `Library "ppx_expect" not found` (opam builds `xla` +before the ppx tree — no dependency edge), then `Library "yojson" not found` / +`int_repr` missing (nothing pulls them in). + +**Why it matters for the lockfile:** `opam lock` computes the closure from *declared* +metadata, so it will **omit `yojson`/`int_repr`** → the lockfile is incomplete → a fresh +`--locked`/CI install fails. Plus a build-ordering race (xla before its ppx). + +**Correct fix** (this is exactly the orc pattern — orc forks+fixes its pinned deps, e.g. +its `jsonaf` "oxcaml-portable-fix" branch): add the missing deps to +`mt-caret/ocaml-xla`'s `dune-project (depends …)` — +`int_repr yojson dune-configurator ppx_expect ppx_sexp_conv ppx_sexp_message ppx_compare` +(or just `ppx_jane` for the ppx set) — regenerate `xla.opam`, push to a branch, and pin +`fox` to that commit. This fixes both presence (lockfile completeness) and ordering. +**This touches a SEPARATE repo (mt-caret/ocaml-xla) — get my OK before pushing.** + +Local-only stopgap that unblocked the build (NOT enough for clean reproducibility): +`opam install yojson int_repr -y` into the switch first, then +`opam install . --deps-only --with-test` (xla then builds because the ppx tree is present). + +## Source porting needed (in order) + +### 1. Effect-handler syntax (5.3 → 5.2) — DONE on the old box, redo it +`lib/core/fox_core.ml` has 4 deep handlers (`Eval`, `Jvp`, `Staging`, `Partial`) using the +OCaml-5.3 syntax `try f () with | effect Fox_effect.Op op, k -> … continue k value`, which +5.2.0+ox can't parse ("Syntax error"). There is a single effect constructor +`Fox_effect.Op : Value0.t Op.t -> Value0.t Effect.t`, and all 4 produce a `Value.t` +(`Value.t = Value0.t`). Add one shared helper and rewrite each handler through it: +```ocaml +(* [Fox_effect.Op] is the only effect performed; one shared deep handler suffices. + [handle] turns an op into the value to resume with. *) +let handle_op ~f ~handle = + Effect.Deep.try_with f () + { effc = + (fun (type a) (eff : a Effect.t) -> + match eff with + | Fox_effect.Op op -> Some (fun k -> continue k (handle op)) + | _ -> None) + } +;; +``` +Each `handle [t] ~f = try f () with | effect Fox_effect.Op op, k -> ; continue k X` +becomes `handle [t] ~f = handle_op ~f ~handle:(fun op -> ; X)` (drop `continue k`, +return the value, close the lambda). `open! Effect` and `open! Effect.Deep` are already at +the top. `fox_effect.ml/.mli` (the `type _ Effect.t += Op …` extension) and `value.ml`'s +`Effect.perform (Fox_effect.Op op)` are fine on 5.2 — no change. + +### 2. `Set_once` API → `[%call_pos]` — DONE on the old box, redo it +Installed sigs (`_opam/lib/core/set_once.mli`): +`set_exn : 'a t -> here:[%call_pos] -> 'a -> unit`, +`get_exn : here:[%call_pos] -> 'a t -> 'a`. In `fox_core.ml`: +- `Set_once.set_exn out_tree_def here (Value_tree.to_def out_tree)` → + `Set_once.set_exn out_tree_def ~here (Value_tree.to_def out_tree)`. +- `Set_once.get_exn out_tree_def [%here]` (2 sites) → `Set_once.get_exn out_tree_def` + (drop `[%here]`; call_pos auto-fills). +(The whole `~here` threading through `flatten_function` could be dropped for `[%call_pos]`, +but the minimal change above works.) + +### 3. OxCaml modes / locality — UNRESOLVED, the hard part +Build stopped here. `lib/core/treeable.ml` `Of_typed_fields.t_of_tree` calls +`Typed_field.create { f = (fun (type a) (field : a Typed_field.t) -> …) }`, but the modal +`Typed_fields_lib.S.create` callback now requires the field at **local** mode: +> Error: This expression has type `'a Typed_field.t -> 'a` but … expected +> `'b Typed_field.t @ local -> 'b`. + +Making `field` local (`fun (type a) (local_ field : a Typed_field.t) -> …`) **cascades**: +the body calls `field_treeable field`, and `field_treeable` is declared in +`Of_typed_fields_arg` (`treeable_intf.ml`) as `'a Typed_field.t -> …` (global). So you'd +also change that module type's `field_treeable` to accept a local field +(`'a Typed_field.t @ local -> …`) and update **every implementation** of +`Of_typed_fields_arg`. + +Next steps that were about to happen: +- Read the exact `create` signature in the installed `typed_fields_lib` `.mli` and check + whether `Typed_field.name`/`get` accept local fields. +- `grep -rn 'typed_fields\|field_treeable\|Of_typed_fields\|@@deriving typed_fields' + lib example` to size the blast radius (the `example/` mnist code likely uses it). +- Expect **more mode errors after this one** (the build stops at the first). Iterate + `dune build` → add `local_` / `@ local` / `@ global` annotations. If it gets deep or + needs real local/global *design* decisions in fox's own API, stop and check with me. + +## Recommended order on the fresh box + +1. Recreate the config scaffolding (above). +2. **Decide the `xla` metadata fix with me first** (push fix to mt-caret/ocaml-xla + pin + the commit) — otherwise the lockfile will be incomplete. +3. Port source: effects (1), Set_once (2), then iterate modes (3) with `dune build`. +4. Download the XLA extension; `opam switch create . 5.2.0+ox --repos ox=…,default + --no-install -y`; then `XLA_EXTENSION_DIR=$PWD/xla_extension opam install . --deps-only + --with-test -y`. +5. `XLA_EXTENSION_DIR=$PWD/xla_extension dune build @default @runtest`; then `dune fmt` + (inspect diff). +6. `opam lock .` → `fox.opam.locked`. **Verify** by doing a clean + `opam switch create 5.2.0+ox --repos ox=…,default --locked` from scratch — this + is where the ocaml-xla metadata bug bites if not fixed. +7. CI (`.github/workflows/build.yml`) + README. Leave changes **uncommitted on main**. + +## CI workflow (drafted; based on orc CI + the ocaml-xla CI) + +Single `ubuntu-latest` job (the xla_extension archive is linux x86_64; macOS needs the +darwin archive). Steps: `jlumbroso/free-disk-space` (the +ox JS tree is large) → +`actions/checkout` → `ocaml/setup-ocaml@v3` with `ocaml-compiler: ocaml-variants.5.2.0+ox` +and `opam-repositories: ox + default` → download xla_extension and +`echo "XLA_EXTENSION_DIR=$PWD/xla_extension" >> "$GITHUB_ENV"` → `actions/cache` on +`~/.opam` + `_opam` keyed on `hashFiles('fox.opam.locked')` → +`opam install . --deps-only --with-test --locked` → `opam exec -- dune build @default` → +`opam exec -- dune runtest` → `opam exec -- dune build @fmt`. + +## README + +Replace the old `DYLD_LIBRARY_PATH=… dune build` line with: note that fox targets OxCaml +(effect syntax) + pins mt-caret/ocaml-xla; one-time setup = download xla_extension + export +`XLA_EXTENSION_DIR` + `opam switch create . 5.2.0+ox --repos ox=…,default --locked`; build/ +test = `XLA_EXTENSION_DIR=$PWD/xla_extension dune build @default @runtest`; regen lockfile = +`opam install . --deps-only --with-test` then `opam lock .`. + +## Open questions to confirm with me + +1. OK to **port fox's source to OxCaml 5.2 + modal libs** (effects + Set_once + the + typed_fields/local-mode changes that touch fox's own `Of_typed_fields_arg` API)? The + alternatives — wait for a 5.3-based OxCaml, or build on vanilla 5.3 + non-modal JS + (which would **not** be OxCaml) — don't meet the stated goal. +2. OK to **push the metadata fix to `mt-caret/ocaml-xla`** and pin fox to that commit? diff --git a/fox.opam b/fox.opam new file mode 100644 index 0000000..7654744 --- /dev/null +++ b/fox.opam @@ -0,0 +1,43 @@ +# This file is generated by dune, edit dune-project instead +opam-version: "2.0" +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." +maintainer: ["mt-caret"] +authors: ["mt-caret"] +homepage: "https://github.com/mt-caret/fox" +bug-reports: "https://github.com/mt-caret/fox/issues" +depends: [ + "dune" {>= "3.20"} + "core" + "core_kernel" + "core_unix" + "base_quickcheck" + "expect_test_helpers_core" + "ppx_jane" + "ppx_typed_fields" + "splittable_random" + "xla" + "ocamlformat" {with-dev-setup} + "odoc" {with-doc | with-dev-setup} + "utop" {with-dev-setup} +] +build: [ + ["dune" "subst"] {dev} + [ + "dune" + "build" + "-p" + name + "-j" + jobs + "@install" + "@runtest" {with-test} + "@doc" {with-doc} + ] +] +dev-repo: "git+https://github.com/mt-caret/fox.git" +pin-depends: [ + ["xla.dev" "git+https://github.com/mt-caret/ocaml-xla.git#6dffd132a9e1b13b3b8c17f1abae110c40ea6f3c"] +] +x-maintenance-intent: ["(latest)"] diff --git a/fox.opam.locked b/fox.opam.locked new file mode 100644 index 0000000..84d370f --- /dev/null +++ b/fox.opam.locked @@ -0,0 +1,242 @@ +opam-version: "2.0" +name: "fox" +version: "dev" +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." +maintainer: "mt-caret" +authors: "mt-caret" +homepage: "https://github.com/mt-caret/fox" +bug-reports: "https://github.com/mt-caret/fox/issues" +depends: [ + "astring" {= "0.8.5"} + "base" {= "v0.18~preview.130.91+190"} + "base-bigarray" {= "base"} + "base-domains" {= "base"} + "base-nnp" {= "base"} + "base-threads" {= "base"} + "base-unix" {= "base"} + "base_bigstring" {= "v0.18~preview.130.91+190"} + "base_quickcheck" {= "v0.18~preview.130.91+190"} + "basement" {= "v0.18~preview.130.91+190"} + "bin_prot" {= "v0.18~preview.130.91+190"} + "camlp-streams" {= "5.0.1"} + "camlzip" {= "1.14"} + "capitalization" {= "v0.18~preview.130.91+190"} + "capsule" {= "v0.18~preview.130.91+190"} + "capsule0" {= "v0.18~preview.130.91+190"} + "cmdliner" {= "1.3.0"} + "conf-autoconf" {= "0.2"} + "conf-libffi" {= "2.0.0"} + "conf-pkg-config" {= "5"} + "conf-which" {= "1"} + "conf-zlib" {= "1"} + "core" {= "v0.18~preview.130.91+190"} + "core_kernel" {= "v0.18~preview.130.91+190"} + "core_unix" {= "v0.18~preview.130.91+190"} + "cppo" {= "1.8.0" & with-doc} + "crunch" {= "4.0.0" & with-doc} + "csexp" {= "1.5.2"} + "ctypes" {= "0.24.0+ox"} + "ctypes-foreign" {= "0.24.0+ox"} + "dune" {= "3.22.2+ox"} + "dune-build-info" {= "3.20.2"} + "dune-configurator" {= "3.22.2+ox"} + "either" {= "1.0.0"} + "expect_test_helpers_core" {= "v0.18~preview.130.91+190"} + "fieldslib" {= "v0.18~preview.130.91+190"} + "fix" {= "20250919"} + "fmt" {= "0.11.0" & with-doc} + "fpath" {= "0.7.3"} + "int_repr" {= "v0.18~preview.130.91+190"} + "integers" {= "0.8.0"} + "jane-street-headers" {= "v0.18~preview.130.91+190"} + "jst-config" {= "v0.18~preview.130.91+190"} + "menhir" {= "20260209"} + "menhirCST" {= "20260209"} + "menhirGLR" {= "20260209"} + "menhirLib" {= "20260209"} + "menhirSdk" {= "20260209"} + "num" {= "1.6"} + "ocaml" {= "5.2.0"} + "ocaml-compiler-libs" {= "v0.17.0+ox"} + "ocaml-config" {= "3"} + "ocaml-options-vanilla" {= "1"} + "ocaml-variants" {= "5.2.0+ox"} + "ocaml-version" {= "4.1.3"} + "ocaml_intrinsics_kernel" {= "v0.18~preview.130.91+190"} + "ocamlbuild" {= "0.16.1+ox"} + "ocamlfind" {= "1.9.8+ox"} + "ocamlformat" {= "0.26.2+ox1"} + "ocamlformat-lib" {= "0.26.2+ox1"} + "ocp-indent" {= "1.9.0"} + "odoc" {= "3.2.1" & with-doc} + "odoc-parser" {= "3.2.1"} + "oxcaml" {= "latest"} + "oxcaml-alcotest" {= "guard"} + "oxcaml-backoff" {= "guard"} + "oxcaml-chrome-trace" {= "guard"} + "oxcaml-cmarkit" {= "guard"} + "oxcaml-compiler" {= "5.2.0minus31"} + "oxcaml-ctypes-foreign-patches" {= "enabled"} + "oxcaml-ctypes-patches" {= "enabled"} + "oxcaml-dot-merlin-reader" {= "guard"} + "oxcaml-dune-action-plugin" {= "guard"} + "oxcaml-dune-build-info" {= "guard"} + "oxcaml-dune-configurator-patches" {= "enabled"} + "oxcaml-dune-glob" {= "guard"} + "oxcaml-dune-patches" {= "enabled"} + "oxcaml-dune-private-libs" {= "guard"} + "oxcaml-dune-rpc" {= "guard"} + "oxcaml-dune-rpc-lwt" {= "guard"} + "oxcaml-dune-site" {= "guard"} + "oxcaml-dyn" {= "guard"} + "oxcaml-eio" {= "guard"} + "oxcaml-eio_linux" {= "guard"} + "oxcaml-eio_main" {= "guard"} + "oxcaml-eio_posix" {= "guard"} + "oxcaml-extlib" {= "guard"} + "oxcaml-fs-io" {= "guard"} + "oxcaml-gen_js_api" {= "guard"} + "oxcaml-js_of_ocaml" {= "guard"} + "oxcaml-js_of_ocaml-compiler" {= "guard"} + "oxcaml-js_of_ocaml-ppx" {= "guard"} + "oxcaml-js_of_ocaml-toplevel" {= "guard"} + "oxcaml-jsonrpc" {= "guard"} + "oxcaml-lsp" {= "guard"} + "oxcaml-lwt" {= "guard"} + "oxcaml-lwt_direct" {= "guard"} + "oxcaml-lwt_ppx" {= "guard"} + "oxcaml-lwt_runtime_events" {= "guard"} + "oxcaml-mdx" {= "guard"} + "oxcaml-merlin" {= "guard"} + "oxcaml-merlin-lib" {= "guard"} + "oxcaml-notty-community" {= "guard"} + "oxcaml-ocaml-compiler-libs-patches" {= "enabled"} + "oxcaml-ocaml-index" {= "guard"} + "oxcaml-ocaml-lsp-server" {= "guard"} + "oxcaml-ocamlbuild-patches" {= "enabled"} + "oxcaml-ocamlc-loc" {= "guard"} + "oxcaml-ocamlfind-patches" {= "enabled"} + "oxcaml-ocamlformat-lib-patches" {= "enabled"} + "oxcaml-ocamlformat-patches" {= "enabled"} + "oxcaml-ojs" {= "guard"} + "oxcaml-omd" {= "guard"} + "oxcaml-opam-core" {= "guard"} + "oxcaml-opam-format" {= "guard"} + "oxcaml-ordering" {= "guard"} + "oxcaml-patch-guards" {= "ox"} + "oxcaml-ppx_deriving" {= "guard"} + "oxcaml-ppxlib-patches" {= "enabled"} + "oxcaml-ppxlib_ast-patches" {= "enabled"} + "oxcaml-re-patches" {= "enabled"} + "oxcaml-sedlex" {= "guard"} + "oxcaml-spawn-patches" {= "enabled"} + "oxcaml-stdune" {= "guard"} + "oxcaml-top-closure" {= "guard"} + "oxcaml-topkg-patches" {= "enabled"} + "oxcaml-utop" {= "guard"} + "oxcaml-uutf-patches" {= "enabled"} + "oxcaml-wasm_of_ocaml-compiler" {= "guard"} + "oxcaml-xdg" {= "guard"} + "oxcaml-yojson-patches" {= "enabled"} + "oxcaml-zarith" {= "guard"} + "parsexp" {= "v0.18~preview.130.91+190"} + "portable" {= "v0.18~preview.130.91+190"} + "ppx_array_base" {= "v0.18~preview.130.91+190"} + "ppx_assert" {= "v0.18~preview.130.91+190"} + "ppx_base" {= "v0.18~preview.130.91+190"} + "ppx_bench" {= "v0.18~preview.130.91+190"} + "ppx_bin_prot" {= "v0.18~preview.130.91+190"} + "ppx_cold" {= "v0.18~preview.130.91+190"} + "ppx_compare" {= "v0.18~preview.130.91+190"} + "ppx_custom_printf" {= "v0.18~preview.130.91+190"} + "ppx_debug_assert" {= "v0.18~preview.130.91+190"} + "ppx_derivers" {= "1.2.1"} + "ppx_diff" {= "v0.18~preview.130.91+190"} + "ppx_disable_unused_warnings" {= "v0.18~preview.130.91+190"} + "ppx_enumerate" {= "v0.18~preview.130.91+190"} + "ppx_expect" {= "v0.18~preview.130.91+190"} + "ppx_fields_conv" {= "v0.18~preview.130.91+190"} + "ppx_fixed_literal" {= "v0.18~preview.130.91+190"} + "ppx_fuelproof" {= "v0.18~preview.130.91+190"} + "ppx_globalize" {= "v0.18~preview.130.91+190"} + "ppx_hash" {= "v0.18~preview.130.91+190"} + "ppx_helpers" {= "v0.18~preview.130.91+190"} + "ppx_here" {= "v0.18~preview.130.91+190"} + "ppx_ignore_instrumentation" {= "v0.18~preview.130.91+190"} + "ppx_inline_test" {= "v0.18~preview.130.91+190"} + "ppx_int63_literal" {= "v0.18~preview.130.91+190"} + "ppx_jane" {= "v0.18~preview.130.91+190"} + "ppx_js_style" {= "v0.18~preview.130.91+190"} + "ppx_let" {= "v0.18~preview.130.91+190"} + "ppx_log" {= "v0.18~preview.130.91+190"} + "ppx_module_timer" {= "v0.18~preview.130.91+190"} + "ppx_optcomp" {= "v0.18~preview.130.91+190"} + "ppx_optional" {= "v0.18~preview.130.91+190"} + "ppx_pipebang" {= "v0.18~preview.130.91+190"} + "ppx_portable" {= "v0.18~preview.130.91+190"} + "ppx_sexp_conv" {= "v0.18~preview.130.91+190"} + "ppx_sexp_message" {= "v0.18~preview.130.91+190"} + "ppx_sexp_value" {= "v0.18~preview.130.91+190"} + "ppx_shorthand" {= "v0.18~preview.130.91+190"} + "ppx_stable" {= "v0.18~preview.130.91+190"} + "ppx_stable_witness" {= "v0.18~preview.130.91+190"} + "ppx_string" {= "v0.18~preview.130.91+190"} + "ppx_string_conv" {= "v0.18~preview.130.91+190"} + "ppx_template" {= "v0.18~preview.130.91+190"} + "ppx_tydi" {= "v0.18~preview.130.91+190"} + "ppx_typed_fields" {= "v0.18~preview.130.91+190"} + "ppx_typerep_conv" {= "v0.18~preview.130.91+190"} + "ppx_var_name" {= "v0.18~preview.130.91+190"} + "ppx_variants_conv" {= "v0.18~preview.130.91+190"} + "ppxlib" {= "0.33.0+ox1"} + "ppxlib_ast" {= "0.33.0+ox1"} + "ppxlib_jane" {= "v0.18~preview.130.91+190"} + "ptime" {= "1.2.0" & with-doc} + "re" {= "1.14.0+ox"} + "result" {= "1.5"} + "seq" {= "base"} + "sexp_pretty" {= "v0.18~preview.130.91+190"} + "sexp_type" {= "v0.18~preview.130.91+190"} + "sexplib" {= "v0.18~preview.130.91+190"} + "sexplib0" {= "v0.18~preview.130.91+190"} + "spawn" {= "v0.15.1+ox"} + "splittable_random" {= "v0.18~preview.130.91+190"} + "stdio" {= "v0.18~preview.130.91+190"} + "stdlib-shims" {= "0.3.0"} + "string_dict" {= "v0.18~preview.130.91+190"} + "time_now" {= "v0.18~preview.130.91+190"} + "topkg" {= "1.1.1+ox"} + "typerep" {= "v0.18~preview.130.91+190"} + "tyxml" {= "4.6.0" & with-doc} + "unique" {= "v0.18~preview.130.91+190"} + "univ_map" {= "v0.18~preview.130.91+190"} + "uopt" {= "v0.18~preview.130.91+190"} + "uucp" {= "17.0.0"} + "uuseg" {= "17.0.0"} + "uutf" {= "1.0.4+ox"} + "variantslib" {= "v0.18~preview.130.91+190"} + "xla" {= "dev"} + "yojson" {= "2.2.2+ox"} +] +build: [ + ["dune" "subst"] {dev} + [ + "dune" + "build" + "-p" + name + "-j" + jobs + "@install" + "@runtest" {with-test} + "@doc" {with-doc} + ] +] +dev-repo: "git+https://github.com/mt-caret/fox.git" +pin-depends: [ + "xla.dev" + "git+https://github.com/mt-caret/ocaml-xla.git#6dffd132a9e1b13b3b8c17f1abae110c40ea6f3c" +] +x-maintenance-intent: ["(latest)"] \ No newline at end of file diff --git a/fox.opam.template b/fox.opam.template new file mode 100644 index 0000000..d52a40c --- /dev/null +++ b/fox.opam.template @@ -0,0 +1,4 @@ +pin-depends: [ + ["xla.dev" "git+https://github.com/mt-caret/ocaml-xla.git#6dffd132a9e1b13b3b8c17f1abae110c40ea6f3c"] +] +x-maintenance-intent: ["(latest)"] diff --git a/lib/core/expr.ml b/lib/core/expr.ml index d4b1d05..9f73a82 100644 --- a/lib/core/expr.ml +++ b/lib/core/expr.ml @@ -1,23 +1,20 @@ open! Core module Var = struct - module T = struct - type t = - { name : string - ; shape : Shape.t - } - [@@deriving compare, sexp, fields ~getters] - end + type t = + { name : string + ; shape : Shape.t + } + [@@deriving compare, hash, sexp, fields ~getters] - include T - include Comparable.Make_plain (T) + include functor Comparable.Make_plain let type_id = Type_equal.Id.create ~name:"Var" [%sexp_of: t] let dims t = shape t |> Shape.dims let to_string { name; shape = { dims; type_ } } = let dims = - Array.to_list dims |> List.map ~f:Int.to_string |> String.concat ~sep:"," + Iarray.to_list dims |> List.map ~f:Int.to_string |> String.concat ~sep:"," in let type_ = [%sexp_of: Type.Packed.t] type_ |> Sexp.to_string |> String.lowercase in [%string "%{name}[%{dims}]: %{type_}"] @@ -25,18 +22,27 @@ module Var = struct end module Atom = struct - type t = + type 'v t = | Var of Var.t - | Value of Value.t - [@@deriving sexp_of] + | Value of 'v + [@@deriving sexp_of, compare, hash] - let to_string = function - | Var { name; shape = { dims = _; type_ = _ } } -> name - | Value value -> Sexp.to_string ([%sexp_of: Value.t] value) + let map t ~f = + match t with + | Var var -> Var var + | Value value -> Value (f value) ;; - let of_value (T { value = x; type_id; shape } as value : Value.t) ~(vars : Var.Set.t) - : t + let to_string t ~value_to_string = + match t with + | Var var -> Var.name var + | Value value -> value_to_string value + ;; + + let of_value + (T { value = x; type_id; shape; id = _ } as value : Value.t) + ~(vars : Var.Set.t) + : Value.t t = match Type_equal.Id.same_witness type_id Var.type_id with | Some T -> @@ -47,7 +53,7 @@ module Atom = struct | None -> Value value ;; - let shape = function + let shape : Value.t t -> Shape.t = function | Var var -> Var.shape var | Value value -> Value.shape value ;; @@ -56,40 +62,68 @@ module Atom = struct end module Eq = struct - type t = + type 'v t = { var : Var.t - ; op : Atom.t Op.t + ; op : 'v Atom.t Op.t } - [@@deriving sexp_of, fields ~getters] + [@@deriving sexp_of, compare, hash, fields ~getters] - let to_string { var; op } = - let op_string = Op.to_string op ~f:Atom.to_string in + let map { var; op } ~f = { var; op = Op.map op ~f:(Atom.map ~f) } + + let to_string { var; op } ~value_to_string = + let op_string = Op.to_string op ~f:(Atom.to_string ~value_to_string) in [%string "%{var#Var} = %{op_string};"] ;; end -type t = +type 'a t = { parameters : Var.t list - ; equations : Eq.t list - ; return_vals : Atom.t Nonempty_list.t + ; consts : 'a Map.M(Var).t + ; equations : 'a Eq.t list + ; return_vals : 'a Atom.t Nonempty_list.t ; out_tree_def : Value_tree.Def.t } -[@@deriving sexp_of, fields ~getters] +[@@deriving sexp_of, compare, hash, fields ~getters] + +let map t ~f = + { t with + consts = Map.map t.consts ~f + ; equations = List.map t.equations ~f:(Eq.map ~f) + ; return_vals = Nonempty_list.map t.return_vals ~f:(Atom.map ~f) + } +;; -let to_string_hum { parameters; equations; return_vals; out_tree_def = _ } = +let to_string_hum + { parameters; consts; equations; return_vals; out_tree_def = _ } + ~value_to_string + = let parameters = String.concat ~sep:" " (List.map parameters ~f:Var.to_string) in - let equations = String.concat ~sep:"\n" (List.map equations ~f:Eq.to_string) in + let consts = + match Map.is_empty consts with + | true -> "" + | false -> + let consts = + Map.to_alist consts + |> List.map ~f:(fun (var, value) -> + [%string " %{var#Var} = %{value_to_string value}"]) + |> String.concat ~sep:"\n" + in + [%string "\nconsts:\n%{consts}"] + in + let equations = + String.concat ~sep:"\n" (List.map equations ~f:(Eq.to_string ~value_to_string)) + in let return_vals = Nonempty_list.to_list return_vals - |> List.map ~f:Atom.to_string + |> List.map ~f:(Atom.to_string ~value_to_string) |> String.concat ~sep:", " in - [%string "%{parameters#String} ->\n%{equations#String}\n( %{return_vals} )"] + [%string "%{parameters#String} ->%{consts}\n%{equations#String}\n( %{return_vals} )"] ;; -let validate ({ parameters; equations; return_vals; out_tree_def = _ } as t) = - let env = Var.Set.of_list parameters in - let validate_atoms ~env (atoms : Atom.t list) = +let validate ({ parameters; consts; equations; return_vals; out_tree_def = _ } as t) = + let env = Var.Set.of_list (parameters @ Map.keys consts) in + let validate_atoms ~env (atoms : Value.t Atom.t list) = match List.filter_map atoms ~f:(function | Var var -> Some var @@ -100,7 +134,9 @@ let validate ({ parameters; equations; return_vals; out_tree_def = _ } as t) = | missing_vars -> raise_s [%message - "Undefined variable" (missing_vars : Var.t list) ~expr:(to_string_hum t)] + "Undefined variable" + (missing_vars : Var.t list) + ~expr:(to_string_hum t ~value_to_string:Value.to_string)] in let env = List.fold equations ~init:env ~f:(fun env { var; op } -> @@ -110,8 +146,8 @@ let validate ({ parameters; equations; return_vals; out_tree_def = _ } as t) = Nonempty_list.to_list return_vals |> validate_atoms ~env ;; -let create ~parameters ~equations ~return_vals ~out_tree_def = - let t = { parameters; equations; return_vals; out_tree_def } in +let create ~parameters ~consts ~equations ~return_vals ~out_tree_def = + let t = { parameters; consts; equations; return_vals; out_tree_def } in validate t; t ;; diff --git a/lib/core/expr.mli b/lib/core/expr.mli index 32c9bd5..37023ef 100644 --- a/lib/core/expr.mli +++ b/lib/core/expr.mli @@ -5,46 +5,49 @@ module Var : sig { name : string ; shape : Shape.t } - [@@deriving compare, sexp, fields ~getters] + [@@deriving compare, hash, sexp, fields ~getters] val type_id : t Type_equal.Id.t - val dims : t -> int array + val dims : t -> int iarray include Comparable.S_plain with type t := t end module Atom : sig - type t = + type 'v t = | Var of Var.t - | Value of Value.t - [@@deriving sexp_of] + | Value of 'v + [@@deriving sexp_of, compare, hash] - val of_value : Value.t -> vars:Var.Set.t -> t - val shape : t -> Shape.t - val dims : t -> int array + val of_value : Value.t -> vars:Var.Set.t -> Value.t t + val shape : Value.t t -> Shape.t + val dims : Value.t t -> int iarray end module Eq : sig - type t = + type 'v t = { var : Var.t - ; op : Atom.t Op.t + ; op : 'v Atom.t Op.t } - [@@deriving sexp_of, fields ~getters] + [@@deriving sexp_of, compare, hash, fields ~getters] end -type t = private +type 'a t = private { parameters : Var.t list - ; equations : Eq.t list - ; return_vals : Atom.t Nonempty_list.t + ; consts : 'a Map.M(Var).t + ; equations : 'a Eq.t list + ; return_vals : 'a Atom.t Nonempty_list.t ; out_tree_def : Value_tree.Def.t } -[@@deriving sexp_of, fields ~getters] +[@@deriving sexp_of, compare, hash, fields ~getters] val create : parameters:Var.t list - -> equations:Eq.t list - -> return_vals:Atom.t Nonempty_list.t + -> consts:Value.t Var.Map.t + -> equations:Value.t Eq.t list + -> return_vals:Value.t Atom.t Nonempty_list.t -> out_tree_def:Value_tree.Def.t - -> t + -> Value.t t -val to_string_hum : t -> string +val map : 'a t -> f:('a -> 'b) -> 'b t +val to_string_hum : 'a t -> value_to_string:('a -> string) -> string diff --git a/lib/core/fox_core.ml b/lib/core/fox_core.ml index 2e7d5e8..c62c1b7 100644 --- a/lib/core/fox_core.ml +++ b/lib/core/fox_core.ml @@ -1,6 +1,4 @@ open! Core -open! Effect -open! Effect.Deep module Expr = Expr module Op = Op module Shape = Shape @@ -9,1076 +7,11 @@ module Treeable = Treeable module Treeable_intf = Treeable_intf module Value = Value module Value_tree = Value_tree +include Handler -let flatten_function - (type in_ out) - (module In : Treeable_intf.S with type t = in_) - (module Out : Treeable_intf.S with type t = out) - ~(f : in_ -> out) - ~in_tree_def - ~here - = - let out_tree_def = Set_once.create () in - ( Staged.stage (fun values -> - let in_tree = Value_tree.unflatten values ~def:in_tree_def in - let out = f (In.t_of_tree in_tree) in - let out_tree = Out.tree_of_t out in - Set_once.set_exn out_tree_def here (Value_tree.to_def out_tree); - Value_tree.flatten out_tree) - , out_tree_def ) -;; +(* Internals exposed only for the test suite. *) +module For_testing = struct + module Partial_value = Partial_eval.Partial_value -module Eval = struct - let handle ~f = - try f () with - | effect Fox_effect.Op op, k -> - let result = - Op.map op ~f:Value.to_tensor_exn - |> Op.eval (module Tensor : Operators_intf.S with type t = Tensor.t) - |> Value.of_tensor - in - continue k result - ;; + let partially_apply_expr_flat = Partial_eval.partially_apply_expr_flat end - -(* foo(x) = x(x + 3) = x^2 + 3x - foo'(x) = 2x + 3 - foo''(x) = 2 - foo'''(x) = 0 -*) -let foo x = Value.O.(x * (x + Value.of_float 3.)) - -let%expect_test "foo" = - Eval.handle ~f:(fun () -> foo (Value.of_float 2.)) |> [%sexp_of: Value.t] |> print_s; - [%expect {| (Tensor 10 Float) |}] -;; - -module Dual_number = struct - type t = - { primal : Value.t - ; tangent : Value.t option - ; id : Id.t - } - [@@deriving sexp_of, fields ~getters] - - let type_id = Type_equal.Id.create ~name:"Dual_number" [%sexp_of: t] - let to_value t : Value.t = T { value = t; type_id; shape = Value.shape t.primal } -end - -module Jvp = struct - type t = { id : Id.t } - - let create () = { id = Id.create () } - - let dual_number t ~primal ~tangent : Dual_number.t = - Option.iter tangent ~f:(fun tangent -> - [%test_eq: int array] (Value.dims primal) (Value.dims tangent); - assert (Type.Packed.equal (Value.type_ primal) (Value.type_ tangent))); - { primal; tangent; id = t.id } - ;; - - let lift t (T { value = _; type_id = _; shape = { dims; type_ } } as value : Value.t) - : Dual_number.t - = - let zeros () = Value.of_typed_tensor (Tensor.Typed.create Float ~dims 0.) in - match Value.coerce value ~type_id:Dual_number.type_id with - | Some x -> - if Id.equal t.id x.id - then x - else dual_number t ~primal:value ~tangent:(Some (zeros ())) - | None -> - (match type_ with - | T Float -> dual_number t ~primal:value ~tangent:(Some (zeros ())) - | T Bool -> dual_number t ~primal:value ~tangent:None) - ;; - - let handle t ~f = - try f () with - | effect Fox_effect.Op op, k -> - let result = - match Op.map op ~f:(lift t) with - | Unary (Neg, a) -> - dual_number - t - ~primal:Value.O.(-a.primal) - ~tangent:(Option.map a.tangent ~f:Value.neg) - | Unary (Sin, a) -> - dual_number - t - ~primal:(Value.sin a.primal) - ~tangent: - (Option.map a.tangent ~f:(fun tangent -> - Value.O.(Value.cos a.primal * tangent))) - | Unary (Cos, a) -> - dual_number - t - ~primal:(Value.cos a.primal) - ~tangent: - (Option.map a.tangent ~f:(fun tangent -> - Value.O.(-Value.sin a.primal * tangent))) - | Unary (Sqrt, a) -> - dual_number - t - ~primal:(Value.sqrt a.primal) - ~tangent: - (Option.map a.tangent ~f:(fun tangent -> - Value.div tangent (Value.scale (Value.sqrt a.primal) 2.))) - | Unary (Exp, a) -> - dual_number - t - ~primal:(Value.exp a.primal) - ~tangent: - (Option.map a.tangent ~f:(fun tangent -> - Value.O.(Value.exp a.primal * tangent))) - | Unary (Log, a) -> - dual_number - t - ~primal:(Value.log a.primal) - ~tangent:(Option.map a.tangent ~f:(fun tangent -> Value.div tangent a.primal)) - | Unary (Sigmoid, a) -> - dual_number - t - ~primal:(Value.sigmoid a.primal) - ~tangent: - (Option.map a.tangent ~f:(fun tangent -> - Value.O.( - Value.sigmoid a.primal - * ((Value.of_float 1. |> Value.broadcast ~dims:(Value.dims a.primal)) - - Value.sigmoid a.primal) - * tangent))) - | Binary (Add, a, b) -> - dual_number - t - ~primal:Value.O.(a.primal + b.primal) - ~tangent: - (Option.map2 a.tangent b.tangent ~f:(fun a_tangent b_tangent -> - Value.O.(a_tangent + b_tangent))) - | Binary (Sub, a, b) -> - dual_number - t - ~primal:Value.O.(a.primal - b.primal) - ~tangent: - (Option.map2 a.tangent b.tangent ~f:(fun a_tangent b_tangent -> - Value.O.(a_tangent - b_tangent))) - | Binary (Mul, a, b) -> - dual_number - t - ~primal:Value.O.(a.primal * b.primal) - ~tangent: - (Option.map2 a.tangent b.tangent ~f:(fun a_tangent b_tangent -> - Value.O.((a_tangent * b.primal) + (a.primal * b_tangent)))) - | Binary (Div, a, b) -> - dual_number - t - ~primal:Value.O.(a.primal / b.primal) - ~tangent: - (Option.map2 a.tangent b.tangent ~f:(fun a_tangent b_tangent -> - Value.O.( - ((a_tangent * b.primal) - (a.primal * b_tangent)) - / (b.primal * b.primal)))) - | Binary (Eq, a, b) -> - dual_number t ~primal:Value.O.(a.primal = b.primal) ~tangent:None - | Binary (Gt, a, b) -> - dual_number t ~primal:Value.O.(a.primal > b.primal) ~tangent:None - | Binary (Lt, a, b) -> - dual_number t ~primal:Value.O.(a.primal < b.primal) ~tangent:None - | Matmul (a, b) -> - dual_number - t - ~primal:(Value.matmul a.primal b.primal) - ~tangent: - (Option.map2 a.tangent b.tangent ~f:(fun a_tangent b_tangent -> - Value.O.( - Value.matmul a_tangent b.primal + Value.matmul a.primal b_tangent))) - | Transpose a -> - dual_number - t - ~primal:(Value.transpose a.primal) - ~tangent:(Option.map a.tangent ~f:Value.transpose) - | Sum { value; dims; keep_dims } -> - dual_number - t - ~primal:(Value.sum ~dims ~keep_dims value.primal) - ~tangent:(Option.map value.tangent ~f:(Value.sum ~dims ~keep_dims)) - | Broadcast { value; dims } -> - dual_number - t - ~primal:(Value.broadcast ~dims value.primal) - ~tangent:(Option.map value.tangent ~f:(Value.broadcast ~dims)) - | Reshape { value; dims } -> - dual_number - t - ~primal:(Value.reshape value.primal ~dims) - ~tangent:(Option.map value.tangent ~f:(Value.reshape ~dims)) - in - continue k (Dual_number.to_value result) - ;; -end - -let jvp - (type in_ out) - (module In : Treeable_intf.S with type t = in_) - (module Out : Treeable_intf.S with type t = out) - ~f - ~(primals : in_) - ~(tangents : in_) - = - let jvp = Jvp.create () in - let primals_tree, tangents_tree = In.tree_of_t primals, In.tree_of_t tangents in - let primals_tree_def, tangents_tree_def = - Value_tree.to_def primals_tree, Value_tree.to_def tangents_tree - in - [%test_eq: Value_tree.Def.t] primals_tree_def tangents_tree_def; - let inputs = - List.zip_exn (Value_tree.flatten primals_tree) (Value_tree.flatten tangents_tree) - |> List.map ~f:(fun (primal, tangent) -> - Dual_number.to_value (Jvp.dual_number jvp ~primal ~tangent:(Some tangent))) - in - let f, out_tree_def = - flatten_function - (module In) - (module Out) - ~f - ~in_tree_def:primals_tree_def - ~here:[%here] - in - let f = Staged.unstage f in - let primals, tangents = - Jvp.handle jvp ~f:(fun () -> f inputs) - |> List.map ~f:(Jvp.lift jvp) - |> List.map ~f:(fun { primal; tangent; id = _ } -> - primal, Option.value_exn ~message:"None tangent not supported in jvp" tangent) - |> List.unzip - in - let out_tree_def = Set_once.get_exn out_tree_def [%here] in - ( Out.t_of_tree (Value_tree.unflatten primals ~def:out_tree_def) - , Out.t_of_tree (Value_tree.unflatten tangents ~def:out_tree_def) ) -;; - -let%expect_test "jvp'" = - Eval.handle ~f:(fun () -> - jvp - (module Value) - (module Value) - ~f:foo - ~primals:(Value.of_float 2.) - ~tangents:(Value.of_float 1.)) - |> [%sexp_of: Value.t * Value.t] - |> print_s; - [%expect {| ((Tensor 10 Float) (Tensor 7 Float)) |}] -;; - -let jvp' ~f ~primal ~tangent = - jvp (module Value) (module Value) ~f ~primals:primal ~tangents:tangent -;; - -let%expect_test "jvp" = - Eval.handle ~f:(fun () -> - jvp' ~f:foo ~primal:(Value.of_float 2.) ~tangent:(Value.of_float 1.)) - |> [%sexp_of: Value.t * Value.t] - |> print_s; - [%expect {| ((Tensor 10 Float) (Tensor 7 Float)) |}]; - Eval.handle ~f:(fun () -> - jvp' - ~f:(fun x -> - let _, tangent = jvp' ~f:foo ~primal:x ~tangent:(Value.of_float 1.) in - tangent) - ~primal:(Value.of_float 2.) - ~tangent:(Value.of_float 1.)) - |> [%sexp_of: Value.t * Value.t] - |> print_s; - [%expect {| ((Tensor 7 Float) (Tensor 2 Float)) |}] -;; - -let derivative ~f ~x = - let (_primal : Value.t), tangent = jvp' ~f ~primal:x ~tangent:(Value.of_float 1.) in - tangent -;; - -let rec nth_order_derivative ~n ~f ~x = - match n with - | 0 -> f x - | _ -> derivative ~f:(fun x -> nth_order_derivative ~n:(n - 1) ~f ~x) ~x -;; - -let%expect_test "nth_order_derivative" = - let print ~n = - Eval.handle ~f:(fun () -> nth_order_derivative ~n ~f:foo ~x:(Value.of_float 2.)) - |> [%sexp_of: Value.t] - |> print_s - in - print ~n:0; - [%expect {| (Tensor 10 Float) |}]; - print ~n:1; - [%expect {| (Tensor 7 Float) |}]; - print ~n:2; - [%expect {| (Tensor 2 Float) |}]; - print ~n:3; - [%expect {| (Tensor 0 Float) |}]; - print ~n:4; - [%expect {| (Tensor 0 Float) |}] -;; - -let%expect_test "pertubation confusion avoidance" = - let f x = - let g (_y : Value.t) = x in - let should_be_zero = derivative ~f:g ~x:(Value.of_float 0.) in - Value.O.(x * should_be_zero) - in - Eval.handle ~f:(fun () -> derivative ~f ~x:(Value.of_float 0.)) - |> [%sexp_of: Value.t] - |> print_s; - [%expect {| (Tensor 0 Float) |}] -;; - -module Staging = struct - type t = - { mutable equations : Expr.Eq.t list - ; mutable name_counter : int - ; mutable vars : Expr.Var.Set.t - } - - let create () = { equations = []; name_counter = 0; vars = Expr.Var.Set.empty } - - let fresh_var t ~shape = - let name = [%string "v_%{t.name_counter#Int}"] in - t.name_counter <- t.name_counter + 1; - let var : Expr.Var.t = { name; shape } in - t.vars <- Set.add t.vars var; - var - ;; - - let value_to_atom t value = Expr.Atom.of_value ~vars:t.vars value - - let handle t ~f = - try f () with - | effect Fox_effect.Op op, k -> - let shape = Op.map op ~f:Value.shape |> Op.infer_shape_exn in - let binder = fresh_var t ~shape in - t.equations <- { var = binder; op = Op.map op ~f:(value_to_atom t) } :: t.equations; - let value : Value.t = T { value = binder; type_id = Expr.Var.type_id; shape } in - continue k value - ;; -end - -(* TODO: could [Expr.t] instead be something like [(in_, out) Expr.t], storing - the modules internally? *) -let build_expr - (type in_ out) - (module In : Treeable_intf.S with type t = in_) - (module Out : Treeable_intf.S with type t = out) - ~f - ~in_tree_def - : Expr.t - = - let staging = Staging.create () in - let parameters = - Value_tree.Def.flatten in_tree_def - |> List.map ~f:(fun dims -> - (* TODO: support arbitrary types here. *) - Staging.fresh_var staging ~shape:{ dims; type_ = T Float }) - in - let f, out_tree_def = - flatten_function (module In) (module Out) ~f ~in_tree_def ~here:[%here] - in - let f = Staged.unstage f in - let result = - Staging.handle staging ~f:(fun () -> - List.map parameters ~f:(fun parameter -> - Value.T - { value = parameter - ; type_id = Expr.Var.type_id - ; shape = Expr.Var.shape parameter - }) - |> f) - in - Expr.create - ~parameters - ~equations:(List.rev staging.equations) - ~return_vals: - (Nonempty_list.of_list_exn result - |> Nonempty_list.map ~f:(Staging.value_to_atom staging)) - ~out_tree_def:(Set_once.get_exn out_tree_def [%here]) -;; - -let build_expr' ~f ~in_dims : Expr.t = - build_expr (module Value) (module Value) ~f ~in_tree_def:(Value.tree_def ~dims:in_dims) -;; - -let%expect_test "build_expr" = - build_expr' ~f:foo ~in_dims:[||] |> Expr.to_string_hum |> print_endline; - [%expect - {| - v_0[]: float -> - v_1[]: float = add v_0 (Tensor 3 Float); - v_2[]: float = mul v_0 v_1; - ( v_2 ) - |}] -;; - -let%expect_test "build_expr2" = - build_expr' ~f:(fun _x -> Value.O.(Value.of_float 2. * Value.of_float 2.)) ~in_dims:[||] - |> Expr.to_string_hum - |> print_endline; - [%expect - {| - v_0[]: float -> - v_1[]: float = mul (Tensor 2 Float) (Tensor 2 Float); - ( v_1 ) - |}] -;; - -let eval_expr_flat (expr : Expr.t) (input : Value.t list) = - let eval_atom (atom : Expr.Atom.t) ~env = - match atom with - | Var var -> Map.find_exn env var - | Value value -> value - in - let env = - List.fold - expr.equations - ~init:(List.zip_exn expr.parameters input |> Expr.Var.Map.of_alist_exn) - ~f:(fun env eq -> - let result = Op.map eq.op ~f:(eval_atom ~env) |> Op.eval (module Value) in - Map.add_exn env ~key:eq.var ~data:result) - in - Nonempty_list.map expr.return_vals ~f:(eval_atom ~env) |> Nonempty_list.to_list -;; - -let eval_expr - (type in_ out) - (module In : Treeable_intf.S with type t = in_) - (module Out : Treeable_intf.S with type t = out) - (expr : Expr.t) - (input : in_) - : out - = - In.tree_of_t input - |> Value_tree.flatten - |> eval_expr_flat expr - |> Value_tree.unflatten ~def:expr.out_tree_def - |> Out.t_of_tree -;; - -let eval_expr' = eval_expr (module Value) (module Value) - -let%expect_test "eval_expr" = - Eval.handle ~f:(fun () -> - eval_expr' (build_expr' ~f:foo ~in_dims:[||]) (Value.of_float 2.)) - |> [%sexp_of: Value.t] - |> print_s; - [%expect {| (Tensor 10 Float) |}] -;; - -let%expect_test "jvp and eval_expr" = - Eval.handle ~f:(fun () -> - jvp' - ~f:(fun x -> eval_expr' (build_expr' ~f:foo ~in_dims:[||]) x) - ~primal:(Value.of_float 2.) - ~tangent:(Value.of_float 1.)) - |> [%sexp_of: Value.t * Value.t] - |> print_s; - [%expect {| ((Tensor 10 Float) (Tensor 7 Float)) |}] -;; - -let%expect_test "nth_order_derivative build_expr" = - let print ~n = - build_expr' ~f:(fun x -> nth_order_derivative ~n ~f:foo ~x) ~in_dims:[||] - |> Expr.to_string_hum - |> print_endline - in - print ~n:0; - [%expect - {| - v_0[]: float -> - v_1[]: float = add v_0 (Tensor 3 Float); - v_2[]: float = mul v_0 v_1; - ( v_2 ) - |}]; - print ~n:1; - [%expect - {| - v_0[]: float -> - v_1[]: float = add (Tensor 1 Float) (Tensor 0 Float); - v_2[]: float = add v_0 (Tensor 3 Float); - v_3[]: float = mul v_0 v_1; - v_4[]: float = mul (Tensor 1 Float) v_2; - v_5[]: float = add v_4 v_3; - v_6[]: float = mul v_0 v_2; - ( v_5 ) - |}]; - print ~n:2; - [%expect - {| - v_0[]: float -> - v_1[]: float = add (Tensor 0 Float) (Tensor 0 Float); - v_2[]: float = add (Tensor 1 Float) (Tensor 0 Float); - v_3[]: float = add (Tensor 1 Float) (Tensor 0 Float); - v_4[]: float = add v_0 (Tensor 3 Float); - v_5[]: float = mul v_0 v_1; - v_6[]: float = mul (Tensor 1 Float) v_2; - v_7[]: float = add v_6 v_5; - v_8[]: float = mul v_0 v_2; - v_9[]: float = mul (Tensor 1 Float) v_3; - v_10[]: float = mul (Tensor 0 Float) v_4; - v_11[]: float = add v_10 v_9; - v_12[]: float = mul (Tensor 1 Float) v_4; - v_13[]: float = add v_11 v_7; - v_14[]: float = add v_12 v_8; - v_15[]: float = mul v_0 v_3; - v_16[]: float = mul (Tensor 1 Float) v_4; - v_17[]: float = add v_16 v_15; - v_18[]: float = mul v_0 v_4; - ( v_13 ) - |}] -;; - -module Partial_value = struct - type t = - | Known of Value.t - | Unknown of Expr.Var.t - [@@deriving sexp_of] - - let shape = function - | Known value -> Value.shape value - | Unknown var -> Expr.Var.shape var - ;; - - let to_atom t ~vars = - match t with - | Known value -> Expr.Atom.of_value value ~vars - | Unknown var -> Expr.Atom.Var var - ;; - - let type_id = Type_equal.Id.create ~name:"Partial_value" [%sexp_of: t] -end - -module Partial = struct - type t = - { mutable equations : Expr.Eq.t list - ; mutable name_counter : int - ; mutable vars : Expr.Var.Set.t - } - - let create () = { equations = []; name_counter = 0; vars = Expr.Var.Set.empty } - - let fresh_var t ~shape = - let name = [%string "p_%{t.name_counter#Int}"] in - t.name_counter <- t.name_counter + 1; - let var : Expr.Var.t = { name; shape } in - t.vars <- Set.add t.vars var; - var - ;; - - let lift value = - Value.coerce value ~type_id:Partial_value.type_id - |> Option.value ~default:(Known value) - ;; - - let value_to_atom t value = Partial_value.to_atom value ~vars:t.vars - - let handle t ~f = - try f () with - | effect Fox_effect.Op op, k -> - let result : Partial_value.t = - match Op.map op ~f:lift with - | Unary (kind, Known a) -> Known (Op.eval (module Value) (Unary (kind, a))) - | Binary (kind, Known a, Known b) -> - Known (Op.eval (module Value) (Binary (kind, a, b))) - | Matmul (Known a, Known b) -> Known (Value.matmul a b) - | Transpose (Known a) -> Known (Value.transpose a) - | Sum { value = Known a; dims; keep_dims } -> Known (Value.sum a ~dims ~keep_dims) - | Broadcast { value = Known a; dims } -> Known (Value.broadcast a ~dims) - | Reshape { value = Known a; dims } -> Known (Value.reshape a ~dims) - | ( Unary ((Neg | Sin | Cos | Sqrt | Exp | Log | Sigmoid), _) - | Binary ((Add | Sub | Mul | Div | Eq | Gt | Lt), _, _) - | Matmul _ | Transpose _ | Sum _ | Broadcast _ | Reshape _ ) as op -> - let shape = Op.map op ~f:Partial_value.shape |> Op.infer_shape_exn in - let binder = fresh_var t ~shape in - t.equations - <- { var = binder; op = Op.map op ~f:(value_to_atom t) } :: t.equations; - Unknown binder - in - continue - k - (T - { value = result - ; type_id = Partial_value.type_id - ; shape = Partial_value.shape result - }) - ;; -end - -(** Do we need the const argument to prevent constants from being instantiated many many times? - - arguably we could just have an eq in the jaxpr for the constant -*) -let partially_apply_expr_flat - (inputs : Partial_value.t list) - ~(f : Value.t list -> Value.t list * Value_tree.Def.t) - : Partial_value.t list * Expr.t - = - let partial = Partial.create () in - let outputs, out_tree_def = - Partial.handle partial ~f:(fun () -> - List.map inputs ~f:(fun input -> - Value.T - { value = input - ; type_id = Partial_value.type_id - ; shape = Partial_value.shape input - }) - |> f) - in - let outputs = List.map outputs ~f:Partial.lift in - let only_unknowns = - List.filter_map ~f:(function - | Partial_value.Known _ -> None - | Unknown var -> Some var) - in - ( outputs - , match - Expr.create - ~parameters:(only_unknowns inputs) - ~equations:(List.rev partial.equations) - ~return_vals: - (only_unknowns outputs - |> List.map ~f:(fun var -> Expr.Atom.Var var) - |> Nonempty_list.of_list_exn) - ~out_tree_def - with - | exception exn -> - raise_s - [%message "Failed to create expr" (exn : exn) (inputs : Partial_value.t list)] - | expr -> expr ) -;; - -let%expect_test "partially_apply_expr_flat" = - let partial_values, expr = - Eval.handle ~f:(fun () -> - partially_apply_expr_flat - [ Known (Value.of_float 2.) - ; Unknown { name = "x"; shape = { dims = [||]; type_ = T Float } } - ] - ~f:(function - | [ x; y ] -> - let x2 = Value.O.(x * x) in - ( [ x2; Value.O.((x2 * y) + Value.of_float 3.); x; Value.O.((y * y) + x2) ] - , Value.tree_def ~dims:[||] ) - | _ -> assert false)) - in - print_s ([%sexp_of: Partial_value.t list] partial_values); - [%expect - {| - ((Known (Tensor 4 Float)) - (Unknown ((name p_3) (shape ((dims ()) (type_ Float))))) - (Known (Tensor 2 Float)) - (Unknown ((name p_1) (shape ((dims ()) (type_ Float)))))) - |}]; - Expr.to_string_hum expr |> print_endline; - [%expect - {| - x[]: float -> - p_0[]: float = mul x x; - p_1[]: float = add p_0 (Tensor 4 Float); - p_2[]: float = mul (Tensor 4 Float) x; - p_3[]: float = add p_2 (Tensor 3 Float); - ( p_3, p_1 ) - |}] -;; - -let linearize - (type in_ out) - (module In : Treeable_intf.S with type t = in_) - (module Out : Treeable_intf.S with type t = out) - ~(f : in_ -> out) - ~(primals : in_) - = - let primals_tree = In.tree_of_t primals in - let primals_tree_def = Value_tree.to_def primals_tree in - let primals = - Value_tree.flatten primals_tree - |> List.map ~f:(fun value -> Partial_value.Known value) - in - let primals_length = List.length primals in - let inputs = - List.append - primals - (List.mapi primals ~f:(fun i primal -> - Partial_value.Unknown - { name = [%string "a_%{i#Int}"]; shape = Partial_value.shape primal })) - in - let outputs, expr = - partially_apply_expr_flat inputs ~f:(fun inputs -> - let primals, tangents = List.split_n inputs primals_length in - let out_primal, out_tangent = - jvp - (module In) - (module Out) - ~f - ~primals:(Value_tree.unflatten primals ~def:primals_tree_def |> In.t_of_tree) - ~tangents:(Value_tree.unflatten tangents ~def:primals_tree_def |> In.t_of_tree) - in - let out_primal_tree = Out.tree_of_t out_primal in - let out_tree_def = Value_tree.to_def out_primal_tree in - let out_tangent_tree = Out.tree_of_t out_tangent in - [%test_eq: Value_tree.Def.t] out_tree_def (Value_tree.to_def out_tangent_tree); - ( List.append - (Value_tree.flatten out_primal_tree) - (Value_tree.flatten out_tangent_tree) - , out_tree_def )) - in - let outputs = List.take outputs (List.length outputs / 2) in - let output = - List.filter_map outputs ~f:(function - | Partial_value.Known value -> Some value - | Unknown _ -> - raise_s - [%message - "unexpected unknown primal" - (outputs : Partial_value.t list) - ~expr:(Expr.to_string_hum expr)]) - |> Value_tree.unflatten ~def:expr.out_tree_def - |> Out.t_of_tree - in - let f_lin (tangents : in_) = - In.tree_of_t tangents - |> Value_tree.flatten - |> eval_expr_flat expr - |> Value_tree.unflatten ~def:expr.out_tree_def - |> Out.t_of_tree - in - output, f_lin -;; - -let linearize' ~f ~primals = linearize (module Value) (module Value) ~f ~primals - -let%expect_test "linearize" = - let y, f_lin = - Eval.handle ~f:(fun () -> linearize' ~f:Value.sin ~primals:(Value.of_float 3.)) - in - print_s [%message "" (y : Value.t) (Float.sin 3. : float)]; - [%expect - {| ((y (Tensor 0.14112000805986721 Float)) ("Float.sin 3." 0.14112000805986721)) |}]; - let y' = Eval.handle ~f:(fun () -> f_lin (Value.of_float 1.)) in - print_s [%message "" (y' : Value.t) (Float.cos 3. : float)]; - [%expect - {| - ((y' (Tensor -0.98999249660044542 Float)) - ("Float.cos 3." -0.98999249660044542)) - |}]; - let y, f_lin = - Eval.handle ~f:(fun () -> - linearize' - ~f:(fun x -> - let y = Value.O.(Value.sin x * Value.of_float 2.) in - Value.O.(-y + x)) - ~primals:(Value.of_float 3.)) - in - let y' = Eval.handle ~f:(fun () -> f_lin (Value.of_float 1.)) in - print_s [%message "" (y : Value.t) (y' : Value.t)]; - [%expect - {| - ((y (Tensor 2.7177599838802657 Float)) - (y' (Tensor 2.9799849932008908 Float))) - |}]; - let f a = - let b = Value.sin a in - let c = Value.neg b in - c - in - build_expr' ~f ~in_dims:[||] |> Expr.to_string_hum |> print_endline; - [%expect - {| - v_0[]: float -> - v_1[]: float = sin v_0; - v_2[]: float = neg v_1; - ( v_2 ) - |}]; - build_expr - (module Value.Tuple2) - (module Value.Tuple2) - ~f:(fun (a, b) -> jvp' ~f ~primal:a ~tangent:b) - ~in_tree_def:(Value.Tuple2.tree_def ~dims1:[||] ~dims2:[||]) - |> Expr.to_string_hum - |> print_endline; - [%expect - {| - v_0[]: float v_1[]: float -> - v_2[]: float = cos v_0; - v_3[]: float = mul v_2 v_1; - v_4[]: float = sin v_0; - v_5[]: float = neg v_3; - v_6[]: float = neg v_4; - ( v_6, v_5 ) - |}]; - (* TODO: fix consts? *) - build_expr' - ~f:(fun x -> - let y, _f_lin = linearize' ~f ~primals:x in - y) - ~in_dims:[||] - |> Expr.to_string_hum - |> print_endline; - [%expect - {| - v_0[]: float -> - v_1[]: float = cos v_0; - v_2[]: float = sin v_0; - v_3[]: float = neg v_2; - ( v_3 ) - |}]; - let _y, f_lin = Eval.handle ~f:(fun () -> linearize' ~f ~primals:(Value.of_float 0.)) in - build_expr' ~f:f_lin ~in_dims:[||] |> Expr.to_string_hum |> print_endline; - [%expect - {| - v_0[]: float -> - v_1[]: float = mul (Tensor 1 Float) v_0; - v_2[]: float = neg v_1; - ( v_2 ) - |}] -;; - -let eval_expr_transposed (expr : Expr.t) args ~cotangents = - let accum_gradient ~ct_env var value = - Map.update ct_env var ~f:(function - | None -> value - | Some existing -> Value.O.(existing + value)) - in - let read_gradient ~ct_env var = - (* TODO: some sort of type inference / add a new variant for "zero"? *) - Map.find ct_env var - |> Option.value_or_thunk ~default:(fun () -> - Tensor.Typed.zeros ~dims:(Expr.Var.dims var) |> Value.of_typed_tensor) - in - let ct_env = - List.zip_exn (Nonempty_list.to_list expr.return_vals) cotangents - |> List.fold ~init:Expr.Var.Map.empty ~f:(fun ct_env (return_val, cotangent) -> - match return_val with - | Value _ -> - (* TODO: do we actually want to just ignore constnats? *) - raise_s [%message "unexpected const return value" (return_val : Expr.Atom.t)] - | Var var -> accum_gradient ~ct_env var cotangent) - in - let ct_env = - List.rev expr.equations - |> List.fold ~init:ct_env ~f:(fun ct_env { var; op } -> - let cotangent = read_gradient ~ct_env var in - let ct_env = - match op with - | Unary (Neg, Var var) -> accum_gradient ~ct_env var (Value.neg cotangent) - | Unary (Sin, Var var) -> accum_gradient ~ct_env var (Value.cos cotangent) - | Unary (Cos, Var var) -> - accum_gradient ~ct_env var (Value.neg (Value.sin cotangent)) - | Unary (Exp, Var var) -> accum_gradient ~ct_env var (Value.exp cotangent) - | Binary (Add, Var var, Value _) | Binary (Add, Value _, Var var) -> - accum_gradient ~ct_env var cotangent - | Binary (Add, Var v1, Var v2) -> - let ct_env = accum_gradient ~ct_env v1 cotangent in - accum_gradient ~ct_env v2 cotangent - | Binary (Sub, Var var, Value _) -> accum_gradient ~ct_env var cotangent - | Binary (Sub, Value _, Var var) -> - accum_gradient ~ct_env var (Value.neg cotangent) - | Binary (Sub, Var v1, Var v2) -> - let ct_env = accum_gradient ~ct_env v1 cotangent in - accum_gradient ~ct_env v2 (Value.neg cotangent) - | Binary (Mul, Var var, Value v) | Binary (Mul, Value v, Var var) -> - accum_gradient ~ct_env var (Value.mul v cotangent) - | Binary (Div, Var var, Value v) -> - accum_gradient ~ct_env var (Value.div cotangent v) - | Matmul (Var var, Value v) -> - accum_gradient ~ct_env var (Value.matmul cotangent (Value.transpose v)) - | Matmul (Value v, Var var) -> - accum_gradient ~ct_env var (Value.matmul (Value.transpose v) cotangent) - | Transpose (Var var) -> accum_gradient ~ct_env var (Value.transpose cotangent) - | Sum { value = Var var; dims; keep_dims } -> - let var_shape = Expr.Var.shape var in - (match keep_dims with - | true -> cotangent - | false -> - (* When dims aren't kept, there are situations where broadcasting to the input - dimension doesn't work e.g. a sum s.t. [ 2; 3 ] -> [ 2 ] *) - let shape_if_dims_were_kept = - Op.infer_shape_exn (Op.Sum { value = var_shape; dims; keep_dims = true }) - in - Value.reshape cotangent ~dims:(Shape.dims shape_if_dims_were_kept)) - |> Value.broadcast ~dims:(Shape.dims var_shape) - |> accum_gradient ~ct_env var - | Broadcast { value = Var var; dims = to_dims } -> - let from_dims = Expr.Var.dims var in - let padding_length = Array.length to_dims - Array.length from_dims in - let non_padded_broadcasts = - Array.sub to_dims ~pos:padding_length ~len:(Array.length from_dims) - |> Array.zip_exn from_dims - |> Array.filter_mapi ~f:(fun i (from, to_) -> - if from <> to_ then Some i else None) - in - let unpadded_cotangent = - match padding_length with - | 0 -> cotangent - | _ -> - Value.sum - cotangent - ~dims:(`Just (Nonempty_list.init padding_length ~f:Fn.id)) - ~keep_dims:false - in - (match Array.to_list non_padded_broadcasts |> Nonempty_list.of_list with - | None -> unpadded_cotangent - | Some non_padded_broadcasts -> - Value.sum - unpadded_cotangent - ~dims:(`Just non_padded_broadcasts) - ~keep_dims:true) - |> accum_gradient ~ct_env var - | Reshape { value = Var var; dims = _ } -> - Value.reshape cotangent ~dims:(Expr.Var.dims var) |> accum_gradient ~ct_env var - | Unary ((Neg | Sin | Cos | Sqrt | Exp | Log | Sigmoid), _) - | Binary ((Add | Sub | Mul | Div | Eq | Gt | Lt), _, _) - | Matmul _ | Transpose _ | Sum _ | Broadcast _ | Reshape _ -> - raise_s - [%message - "Invalid var/val op combination" - (op : Expr.Atom.t Op.t) - ~expr:(Expr.to_string_hum expr)] - in - ct_env) - in - List.map args ~f:(read_gradient ~ct_env) -;; - -let vjp - (type in_ out) - (module In : Treeable_intf.S with type t = in_) - (module Out : Treeable_intf.S with type t = out) - ~(f : in_ -> out) - ~(primals : in_) - = - let primals_tree = In.tree_of_t primals in - let primals_tree_def = Value_tree.to_def primals_tree in - let primals = - Value_tree.flatten primals_tree - |> List.map ~f:(fun value -> Partial_value.Known value) - in - let primals_length = List.length primals in - let tangent_vars = - List.mapi primals ~f:(fun i primal -> - { Expr.Var.name = [%string "a_%{i#Int}"]; shape = Partial_value.shape primal }) - in - let inputs = - List.append primals (List.map tangent_vars ~f:(fun var -> Partial_value.Unknown var)) - in - let outputs, expr = - partially_apply_expr_flat inputs ~f:(fun inputs -> - let primals, tangents = List.split_n inputs primals_length in - let out_primal, out_tangent = - jvp - (module In) - (module Out) - ~f - ~primals:(Value_tree.unflatten primals ~def:primals_tree_def |> In.t_of_tree) - ~tangents:(Value_tree.unflatten tangents ~def:primals_tree_def |> In.t_of_tree) - in - let out_primal_tree = Out.tree_of_t out_primal in - let out_tree_def = Value_tree.to_def out_primal_tree in - let out_tangent_tree = Out.tree_of_t out_tangent in - [%test_eq: Value_tree.Def.t] out_tree_def (Value_tree.to_def out_tangent_tree); - ( List.append - (Value_tree.flatten out_primal_tree) - (Value_tree.flatten out_tangent_tree) - , out_tree_def )) - in - let outputs = List.take outputs (List.length outputs / 2) in - let output = - List.filter_map outputs ~f:(function - | Partial_value.Known value -> Some value - | Unknown _ -> - raise_s - [%message - "unexpected unknown primal" - (outputs : Partial_value.t list) - ~expr:(Expr.to_string_hum expr)]) - |> Value_tree.unflatten ~def:expr.out_tree_def - |> Out.t_of_tree - in - let f_vjp (cotangents : out) = - match - let cotangents_tree = Out.tree_of_t cotangents in - [%test_result: Value_tree.Def.t] - (Value_tree.to_def cotangents_tree) - ~expect:expr.out_tree_def; - eval_expr_transposed - expr - tangent_vars - ~cotangents:(Value_tree.flatten cotangents_tree) - |> Value_tree.unflatten ~def:primals_tree_def - |> In.t_of_tree - with - | in_ -> in_ - | exception exn -> - Exn.reraise - exn - (Sexp.to_string_hum [%message (exn : exn) ~expr:(Expr.to_string_hum expr)]) - in - output, f_vjp -;; - -let vjp' ~f ~primal = vjp (module Value) (module Value) ~f ~primals:primal - -let grad_and_value - (type in_) - (module In : Treeable_intf.S with type t = in_) - ~(f : in_ -> Value.t) - ~x - = - let y, f_vjp = vjp (module In) (module Value) ~f ~primals:x in - y, f_vjp (Value.of_float 1.) -;; - -let grad_and_value' ~f ~x = grad_and_value (module Value) ~f ~x - -let grad module_ ~f ~x = - let _y, grad = grad_and_value module_ ~f ~x in - grad -;; - -let grad' ~f ~x = grad (module Value) ~f ~x - -let%expect_test "grad" = - let y, f_vjp = - Eval.handle ~f:(fun () -> vjp' ~f:Value.sin ~primal:(Value.of_float 3.)) - in - let y' = Eval.handle ~f:(fun () -> f_vjp (Value.of_float 1.)) in - print_s [%message "" (y : Value.t) (y' : Value.t)]; - [%expect - {| - ((y (Tensor 0.14112000805986721 Float)) - (y' (Tensor -0.98999249660044542 Float))) - |}]; - Eval.handle ~f:(fun () -> grad' ~f:(fun x -> Value.O.(x * x)) ~x:(Value.of_float 3.)) - |> [%sexp_of: Value.t] - |> print_s; - [%expect {| (Tensor 6 Float) |}]; - Eval.handle ~f:(fun () -> - grad' - ~f:(fun x -> - let y = Value.O.(Value.sin x * Value.of_float 2.) in - Value.O.(-y + x)) - ~x:(Value.of_float 3.)) - |> [%sexp_of: Value.t] - |> print_s; - [%expect {| (Tensor 2.9799849932008908 Float) |}]; - Eval.handle ~f:(fun () -> - grad' - ~f:(Value.sum ~keep_dims:false) - ~x:(Value.of_tensor (Tensor.of_list2_exn Float [ [ 1.; 2. ]; [ 3.; 4. ] ]))) - |> [%sexp_of: Value.t] - |> print_s; - [%expect {| (Tensor ((1 1) (1 1)) (dims (2 2)) (type_ Float)) |}]; - Eval.handle ~f:(fun () -> - grad' - ~f:(fun x -> - Value.broadcast x ~dims:[| 3; 4 |] |> Value.sum ~dims:(`Just [ 1 ]) |> Value.mean) - ~x:(Value.of_typed_tensor (Tensor.Typed.arange 4))) - |> [%sexp_of: Value.t] - |> print_s; - [%expect {| (Tensor (1 1 1 1) (dims (4)) (type_ Float)) |}] -;; diff --git a/lib/core/fox_effect.ml b/lib/core/fox_effect.ml index 53a1c6c..aaa420f 100644 --- a/lib/core/fox_effect.ml +++ b/lib/core/fox_effect.ml @@ -3,3 +3,13 @@ open! Effect open! Effect.Deep type _ Effect.t += Op : Value0.t Op.t -> Value0.t t + +let handle ~f ~(handle : Value0.t Op.t -> Value0.t) = + let effc : type a. a Effect.t -> ((a, _) continuation -> _) option = + fun eff -> + match eff with + | Op op -> Some (fun k -> continue k (handle op)) + | _ -> None + in + Effect.Deep.try_with f () { effc } +;; diff --git a/lib/core/fox_effect.mli b/lib/core/fox_effect.mli index 53a1c6c..12a9d30 100644 --- a/lib/core/fox_effect.mli +++ b/lib/core/fox_effect.mli @@ -3,3 +3,7 @@ open! Effect open! Effect.Deep type _ Effect.t += Op : Value0.t Op.t -> Value0.t t + +(** Runs [f], interpreting each [Op] effect it performs with [handle] to produce the value + the computation resumes with. *) +val handle : f:(unit -> 'a) -> handle:(Value0.t Op.t -> Value0.t) -> 'a diff --git a/lib/core/handler.ml b/lib/core/handler.ml new file mode 100644 index 0000000..fada16c --- /dev/null +++ b/lib/core/handler.ml @@ -0,0 +1,442 @@ +open! Core +module Partial_value = Partial_eval.Partial_value + +let partially_apply_expr_flat = Partial_eval.partially_apply_expr_flat + +let flatten_function + (type in_ out) + (module In : Treeable_intf.S with type t = in_) + (module Out : Treeable_intf.S with type t = out) + ~(f : in_ -> out) + ~in_tree_def + ~here + = + let out_tree_def = Set_once.create () in + ( Staged.stage (fun values -> + let in_tree = Value_tree.unflatten values ~def:in_tree_def in + let out = f (In.t_of_tree in_tree) in + let out_tree = Out.tree_of_t out in + Set_once.set_exn out_tree_def ~here (Value_tree.to_def out_tree); + Value_tree.flatten out_tree) + , out_tree_def ) +;; + +let eval ~f = + Fox_effect.handle ~f ~handle:(fun op -> + Op.map op ~f:Value.to_tensor_exn + |> Op.eval (module Tensor : Operators_intf.S with type t = Tensor.t) + |> Value.of_tensor) +;; + +let jvp + (type in_ out) + (module In : Treeable_intf.S with type t = in_) + (module Out : Treeable_intf.S with type t = out) + ~f + ~(primals : in_) + ~(tangents : in_) + = + let jvp = Jvp.create () in + let primals_tree, tangents_tree = In.tree_of_t primals, In.tree_of_t tangents in + let primals_tree_def, tangents_tree_def = + Value_tree.to_def primals_tree, Value_tree.to_def tangents_tree + in + [%test_eq: Value_tree.Def.t] primals_tree_def tangents_tree_def; + let inputs = + List.zip_exn (Value_tree.flatten primals_tree) (Value_tree.flatten tangents_tree) + |> List.map ~f:(fun (primal, tangent) -> + Jvp.Dual_number.to_value (Jvp.dual_number jvp ~primal ~tangent:(Some tangent))) + in + let f, out_tree_def = + flatten_function + (module In) + (module Out) + ~f + ~in_tree_def:primals_tree_def + ~here:[%here] + in + let f = Staged.unstage f in + let primals, tangents = + Jvp.handle jvp ~f:(fun () -> f inputs) + |> List.map ~f:(Jvp.lift jvp) + |> List.map ~f:(fun dual_number -> + ( Jvp.Dual_number.primal dual_number + , Option.value_exn + ~message:"None tangent not supported in jvp" + (Jvp.Dual_number.tangent dual_number) )) + |> List.unzip + in + let out_tree_def = Set_once.get_exn out_tree_def in + ( Out.t_of_tree (Value_tree.unflatten primals ~def:out_tree_def) + , Out.t_of_tree (Value_tree.unflatten tangents ~def:out_tree_def) ) +;; + +let jvp' ~f ~primal ~tangent = + jvp (module Value) (module Value) ~f ~primals:primal ~tangents:tangent +;; + +let derivative ~f ~x = + let (_primal : Value.t), tangent = jvp' ~f ~primal:x ~tangent:(Value.of_float 1.) in + tangent +;; + +let rec nth_order_derivative ~n ~f ~x = + match n with + | 0 -> f x + | _ -> derivative ~f:(fun x -> nth_order_derivative ~n:(n - 1) ~f ~x) ~x +;; + +(* TODO: could [Expr.t] instead be something like [(in_, out) Expr.t], storing the modules + internally? *) +let build_expr + (type in_ out) + (module In : Treeable_intf.S with type t = in_) + (module Out : Treeable_intf.S with type t = out) + ~f + ~in_tree_def + = + let staging = Staging.create () in + let parameters = + Value_tree.Def.flatten in_tree_def + |> List.map ~f:(fun dims -> + (* TODO: support arbitrary types here. *) + Staging.fresh_var staging ~shape:{ dims; type_ = T Float }) + in + let f, out_tree_def = + flatten_function (module In) (module Out) ~f ~in_tree_def ~here:[%here] + in + let f = Staged.unstage f in + let result = + Staging.handle staging ~f:(fun () -> + List.map parameters ~f:(fun parameter -> + Value.create + ~value:parameter + ~type_id:Expr.Var.type_id + ~shape:(Expr.Var.shape parameter)) + |> f) + in + let return_vals = + Nonempty_list.of_list_exn result + |> Nonempty_list.map ~f:(Staging.intern_value staging) + in + Expr.create + ~parameters + ~consts:(Staging.consts_map staging) + ~equations:(Staging.equations staging) + ~return_vals + ~out_tree_def:(Set_once.get_exn out_tree_def) +;; + +let build_expr' ~f ~in_dims = + build_expr (module Value) (module Value) ~f ~in_tree_def:(Value.tree_def ~dims:in_dims) +;; + +let eval_expr_flat (expr : Value.t Expr.t) (input : Value.t list) = + let eval_atom (atom : Value.t Expr.Atom.t) ~env = + match atom with + | Var var -> Map.find_exn env var + | Value value -> value + in + let env = + List.zip_exn expr.parameters input + |> Expr.Var.Map.of_alist_exn + |> Map.merge_disjoint_exn expr.consts + in + let env = + List.fold expr.equations ~init:env ~f:(fun env eq -> + let result = Op.map eq.op ~f:(eval_atom ~env) |> Op.eval (module Value) in + Map.add_exn env ~key:eq.var ~data:result) + in + Nonempty_list.map expr.return_vals ~f:(eval_atom ~env) |> Nonempty_list.to_list +;; + +let eval_expr + (type in_ out) + (module In : Treeable_intf.S with type t = in_) + (module Out : Treeable_intf.S with type t = out) + (expr : Value.t Expr.t) + (input : in_) + : out + = + In.tree_of_t input + |> Value_tree.flatten + |> eval_expr_flat expr + |> Value_tree.unflatten ~def:expr.out_tree_def + |> Out.t_of_tree +;; + +let eval_expr' = eval_expr (module Value) (module Value) + +let linearize + (type in_ out) + (module In : Treeable_intf.S with type t = in_) + (module Out : Treeable_intf.S with type t = out) + ~(f : in_ -> out) + ~(primals : in_) + = + let primals_tree = In.tree_of_t primals in + let primals_tree_def = Value_tree.to_def primals_tree in + let primals = + Value_tree.flatten primals_tree + |> List.map ~f:(fun value -> Partial_value.Known value) + in + let primals_length = List.length primals in + let inputs = + List.append + primals + (List.mapi primals ~f:(fun i primal -> + Partial_value.Unknown + { name = [%string "a_%{i#Int}"]; shape = Partial_value.shape primal })) + in + let outputs, expr = + partially_apply_expr_flat inputs ~f:(fun inputs -> + let primals, tangents = List.split_n inputs primals_length in + let out_primal, out_tangent = + jvp + (module In) + (module Out) + ~f + ~primals:(Value_tree.unflatten primals ~def:primals_tree_def |> In.t_of_tree) + ~tangents:(Value_tree.unflatten tangents ~def:primals_tree_def |> In.t_of_tree) + in + let out_primal_tree = Out.tree_of_t out_primal in + let out_tree_def = Value_tree.to_def out_primal_tree in + let out_tangent_tree = Out.tree_of_t out_tangent in + [%test_eq: Value_tree.Def.t] out_tree_def (Value_tree.to_def out_tangent_tree); + ( List.append + (Value_tree.flatten out_primal_tree) + (Value_tree.flatten out_tangent_tree) + , out_tree_def )) + in + let outputs = List.take outputs (List.length outputs / 2) in + let output = + List.filter_map outputs ~f:(function + | Partial_value.Known value -> Some value + | Unknown _ -> + raise_s + [%message + "unexpected unknown primal" + (outputs : Partial_value.t list) + ~expr:(Expr.to_string_hum expr ~value_to_string:Value.to_string)]) + |> Value_tree.unflatten ~def:expr.out_tree_def + |> Out.t_of_tree + in + let f_lin (tangents : in_) = + In.tree_of_t tangents + |> Value_tree.flatten + |> eval_expr_flat expr + |> Value_tree.unflatten ~def:expr.out_tree_def + |> Out.t_of_tree + in + output, f_lin +;; + +let linearize' ~f ~primals = linearize (module Value) (module Value) ~f ~primals + +let eval_expr_transposed (expr : Value.t Expr.t) args ~cotangents = + let accum_gradient ~ct_env var value = + Map.update ct_env var ~f:(function + | None -> value + | Some existing -> Value.O.(existing + value)) + in + let read_gradient ~ct_env var = + (* TODO: some sort of type inference / add a new variant for "zero"? *) + Map.find ct_env var + |> Option.value_or_thunk ~default:(fun () -> + Tensor.Typed.zeros ~dims:(Expr.Var.dims var) |> Value.of_typed_tensor) + in + let ct_env = + List.zip_exn (Nonempty_list.to_list expr.return_vals) cotangents + |> List.fold ~init:Expr.Var.Map.empty ~f:(fun ct_env (return_val, cotangent) -> + match return_val with + | Value _ -> + (* TODO: do we actually want to just ignore constnats? *) + raise_s + [%message "unexpected const return value" (return_val : Value.t Expr.Atom.t)] + | Var var -> accum_gradient ~ct_env var cotangent) + in + let ct_env = + List.rev expr.equations + |> List.fold ~init:ct_env ~f:(fun ct_env { var; op } -> + let cotangent = read_gradient ~ct_env var in + let ct_env = + match op with + | Unary (Neg, Var var) -> accum_gradient ~ct_env var (Value.neg cotangent) + | Unary (Sin, Var var) -> accum_gradient ~ct_env var (Value.cos cotangent) + | Unary (Cos, Var var) -> + accum_gradient ~ct_env var (Value.neg (Value.sin cotangent)) + | Unary (Exp, Var var) -> accum_gradient ~ct_env var (Value.exp cotangent) + | Binary (Add, Var var, Value _) | Binary (Add, Value _, Var var) -> + accum_gradient ~ct_env var cotangent + | Binary (Add, Var v1, Var v2) -> + let ct_env = accum_gradient ~ct_env v1 cotangent in + accum_gradient ~ct_env v2 cotangent + | Binary (Sub, Var var, Value _) -> accum_gradient ~ct_env var cotangent + | Binary (Sub, Value _, Var var) -> + accum_gradient ~ct_env var (Value.neg cotangent) + | Binary (Sub, Var v1, Var v2) -> + let ct_env = accum_gradient ~ct_env v1 cotangent in + accum_gradient ~ct_env v2 (Value.neg cotangent) + | Binary (Mul, Var var, Value v) | Binary (Mul, Value v, Var var) -> + accum_gradient ~ct_env var (Value.mul v cotangent) + | Binary (Div, Var var, Value v) -> + accum_gradient ~ct_env var (Value.div cotangent v) + | Matmul (Var var, Value v) -> + accum_gradient ~ct_env var (Value.matmul cotangent (Value.transpose v)) + | Matmul (Value v, Var var) -> + accum_gradient ~ct_env var (Value.matmul (Value.transpose v) cotangent) + | Transpose (Var var) -> accum_gradient ~ct_env var (Value.transpose cotangent) + | Sum { value = Var var; dims; keep_dims } -> + let var_shape = Expr.Var.shape var in + (match keep_dims with + | true -> cotangent + | false -> + (* When dims aren't kept, there are situations where broadcasting to the + input dimension doesn't work e.g. a sum s.t. [ 2; 3 ] -> [ 2 ] *) + let shape_if_dims_were_kept = + Op.infer_shape_exn (Op.Sum { value = var_shape; dims; keep_dims = true }) + in + Value.reshape cotangent ~dims:(Shape.dims shape_if_dims_were_kept)) + |> Value.broadcast ~dims:(Shape.dims var_shape) + |> accum_gradient ~ct_env var + | Broadcast { value = Var var; dims = to_dims } -> + let from_dims = Expr.Var.dims var in + let padding_length = Iarray.length to_dims - Iarray.length from_dims in + let non_padded_broadcasts = + Iarray.sub to_dims ~pos:padding_length ~len:(Iarray.length from_dims) + |> Iarray.zip_exn from_dims + |> Iarray.to_list + |> List.filter_mapi ~f:(fun i (from, to_) -> + if from <> to_ then Some i else None) + in + let unpadded_cotangent = + match padding_length with + | 0 -> cotangent + | _ -> + Value.sum + cotangent + ~dims:(`Just (Nonempty_list.init padding_length ~f:Fn.id)) + ~keep_dims:false + in + (match Nonempty_list.of_list non_padded_broadcasts with + | None -> unpadded_cotangent + | Some non_padded_broadcasts -> + Value.sum + unpadded_cotangent + ~dims:(`Just non_padded_broadcasts) + ~keep_dims:true) + |> accum_gradient ~ct_env var + | Reshape { value = Var var; dims = _ } -> + Value.reshape cotangent ~dims:(Expr.Var.dims var) |> accum_gradient ~ct_env var + | Unary ((Neg | Sin | Cos | Sqrt | Exp | Log | Sigmoid), _) + | Binary ((Add | Sub | Mul | Div | Eq | Gt | Lt), _, _) + | Matmul _ | Transpose _ | Sum _ | Broadcast _ | Reshape _ -> + raise_s + [%message + "Invalid var/val op combination" + (op : Value.t Expr.Atom.t Op.t) + ~expr:(Expr.to_string_hum expr ~value_to_string:Value.to_string)] + in + ct_env) + in + List.map args ~f:(read_gradient ~ct_env) +;; + +let vjp + (type in_ out) + (module In : Treeable_intf.S with type t = in_) + (module Out : Treeable_intf.S with type t = out) + ~(f : in_ -> out) + ~(primals : in_) + = + let primals_tree = In.tree_of_t primals in + let primals_tree_def = Value_tree.to_def primals_tree in + let primals = + Value_tree.flatten primals_tree + |> List.map ~f:(fun value -> Partial_value.Known value) + in + let primals_length = List.length primals in + let tangent_vars = + List.mapi primals ~f:(fun i primal -> + { Expr.Var.name = [%string "a_%{i#Int}"]; shape = Partial_value.shape primal }) + in + let inputs = + List.append primals (List.map tangent_vars ~f:(fun var -> Partial_value.Unknown var)) + in + let outputs, expr = + partially_apply_expr_flat inputs ~f:(fun inputs -> + let primals, tangents = List.split_n inputs primals_length in + let out_primal, out_tangent = + jvp + (module In) + (module Out) + ~f + ~primals:(Value_tree.unflatten primals ~def:primals_tree_def |> In.t_of_tree) + ~tangents:(Value_tree.unflatten tangents ~def:primals_tree_def |> In.t_of_tree) + in + let out_primal_tree = Out.tree_of_t out_primal in + let out_tree_def = Value_tree.to_def out_primal_tree in + let out_tangent_tree = Out.tree_of_t out_tangent in + [%test_eq: Value_tree.Def.t] out_tree_def (Value_tree.to_def out_tangent_tree); + ( List.append + (Value_tree.flatten out_primal_tree) + (Value_tree.flatten out_tangent_tree) + , out_tree_def )) + in + let outputs = List.take outputs (List.length outputs / 2) in + let output = + List.filter_map outputs ~f:(function + | Partial_value.Known value -> Some value + | Unknown _ -> + raise_s + [%message + "unexpected unknown primal" + (outputs : Partial_value.t list) + ~expr:(Expr.to_string_hum expr ~value_to_string:Value.to_string)]) + |> Value_tree.unflatten ~def:expr.out_tree_def + |> Out.t_of_tree + in + let f_vjp (cotangents : out) = + match + let cotangents_tree = Out.tree_of_t cotangents in + [%test_result: Value_tree.Def.t] + (Value_tree.to_def cotangents_tree) + ~expect:expr.out_tree_def; + eval_expr_transposed + expr + tangent_vars + ~cotangents:(Value_tree.flatten cotangents_tree) + |> Value_tree.unflatten ~def:primals_tree_def + |> In.t_of_tree + with + | in_ -> in_ + | exception exn -> + Exn.reraise + exn + (Sexp.to_string_hum + [%message + (exn : exn) ~expr:(Expr.to_string_hum expr ~value_to_string:Value.to_string)]) + in + output, f_vjp +;; + +let vjp' ~f ~primal = vjp (module Value) (module Value) ~f ~primals:primal + +let grad_and_value + (type in_) + (module In : Treeable_intf.S with type t = in_) + ~(f : in_ -> Value.t) + ~x + = + let y, f_vjp = vjp (module In) (module Value) ~f ~primals:x in + y, f_vjp (Value.of_float 1.) +;; + +let grad_and_value' ~f ~x = grad_and_value (module Value) ~f ~x + +let grad module_ ~f ~x = + let _y, grad = grad_and_value module_ ~f ~x in + grad +;; + +let grad' ~f ~x = grad (module Value) ~f ~x diff --git a/lib/core/handler.mli b/lib/core/handler.mli new file mode 100644 index 0000000..5e48e13 --- /dev/null +++ b/lib/core/handler.mli @@ -0,0 +1,86 @@ +open! Core + +(** Runs [f], evaluating each tensor op eagerly on the [Tensor] backend. *) +val eval : f:(unit -> 'a) -> 'a + +(** Forward-mode AD: evaluates [f] at [primals] and its directional derivative along + [tangents], returning both the primal and tangent outputs. *) +val jvp + : (module Treeable_intf.S with type t = 'in_) + -> (module Treeable_intf.S with type t = 'out) + -> f:('in_ -> 'out) + -> primals:'in_ + -> tangents:'in_ + -> 'out * 'out + +val jvp' + : f:(Value.t -> Value.t) + -> primal:Value.t + -> tangent:Value.t + -> Value.t * Value.t + +(** The derivative of a scalar [f] at [x] (its tangent for a unit input tangent). *) +val derivative : f:(Value.t -> Value.t) -> x:Value.t -> Value.t + +val nth_order_derivative : n:int -> f:(Value.t -> Value.t) -> x:Value.t -> Value.t + +(** Traces [f] into an [Expr.t]; constants it closes over are hoisted into [consts]. *) +val build_expr + : (module Treeable_intf.S with type t = 'in_) + -> (module Treeable_intf.S with type t = 'out) + -> f:('in_ -> 'out) + -> in_tree_def:Value_tree.Def.t + -> Value.t Expr.t + +val build_expr' : f:(Value.t -> Value.t) -> in_dims:int iarray -> Value.t Expr.t + +(** Evaluates a traced [Expr.t] on the surrounding effect handler. *) +val eval_expr + : (module Treeable_intf.S with type t = 'in_) + -> (module Treeable_intf.S with type t = 'out) + -> Value.t Expr.t + -> 'in_ + -> 'out + +val eval_expr' : Value.t Expr.t -> Value.t -> Value.t + +(** Linearizes [f] at [primals]: returns the primal output and the linear tangent map. *) +val linearize + : (module Treeable_intf.S with type t = 'in_) + -> (module Treeable_intf.S with type t = 'out) + -> f:('in_ -> 'out) + -> primals:'in_ + -> 'out * ('in_ -> 'out) + +val linearize' + : f:(Value.t -> Value.t) + -> primals:Value.t + -> Value.t * (Value.t -> Value.t) + +(** Reverse-mode AD: returns the primal output and a function mapping an output cotangent + to the corresponding input cotangent. *) +val vjp + : (module Treeable_intf.S with type t = 'in_) + -> (module Treeable_intf.S with type t = 'out) + -> f:('in_ -> 'out) + -> primals:'in_ + -> 'out * ('out -> 'in_) + +val vjp' : f:(Value.t -> Value.t) -> primal:Value.t -> Value.t * (Value.t -> Value.t) + +(** [grad]ient of a scalar-valued [f], paired with its value. *) +val grad_and_value + : (module Treeable_intf.S with type t = 'in_) + -> f:('in_ -> Value.t) + -> x:'in_ + -> Value.t * 'in_ + +val grad_and_value' : f:(Value.t -> Value.t) -> x:Value.t -> Value.t * Value.t + +val grad + : (module Treeable_intf.S with type t = 'in_) + -> f:('in_ -> Value.t) + -> x:'in_ + -> 'in_ + +val grad' : f:(Value.t -> Value.t) -> x:Value.t -> Value.t diff --git a/lib/core/id.ml b/lib/core/id.ml index 3df799a..cb9ffb7 100644 --- a/lib/core/id.ml +++ b/lib/core/id.ml @@ -1,6 +1,6 @@ open! Core -type t = int [@@deriving equal, sexp_of] +type t = int [@@deriving compare, equal, sexp_of] let create = let counter = ref 0 in diff --git a/lib/core/id.mli b/lib/core/id.mli index 8fe9fc4..324c133 100644 --- a/lib/core/id.mli +++ b/lib/core/id.mli @@ -1,7 +1,7 @@ -(** An opaque unique identifier.*) +(** An opaque unique identifier. *) open! Core -type t = private int [@@deriving equal, sexp_of] +type t = private int [@@deriving compare, equal, sexp_of] val create : unit -> t diff --git a/lib/core/jvp.ml b/lib/core/jvp.ml new file mode 100644 index 0000000..d8fd84e --- /dev/null +++ b/lib/core/jvp.ml @@ -0,0 +1,160 @@ +open! Core + +module Dual_number = struct + type t = + { primal : Value.t + ; tangent : Value.t option + ; id : Id.t + } + [@@deriving sexp_of, fields ~getters] + + let type_id = Type_equal.Id.create ~name:"Dual_number" [%sexp_of: t] + let to_value t : Value.t = Value.create ~value:t ~type_id ~shape:(Value.shape t.primal) +end + +type t = { id : Id.t } + +let create () = { id = Id.create () } + +let dual_number t ~primal ~tangent : Dual_number.t = + Option.iter tangent ~f:(fun tangent -> + [%test_eq: int iarray] (Value.dims primal) (Value.dims tangent); + assert (Type.Packed.equal (Value.type_ primal) (Value.type_ tangent))); + { primal; tangent; id = t.id } +;; + +let lift + t + (T { value = _; type_id = _; shape = { dims; type_ }; id = _ } as value : Value.t) + : Dual_number.t + = + let zeros () = Value.of_typed_tensor (Tensor.Typed.create Float ~dims 0.) in + match Value.coerce value ~type_id:Dual_number.type_id with + | Some x -> + if Id.equal t.id x.id + then x + else dual_number t ~primal:value ~tangent:(Some (zeros ())) + | None -> + (match type_ with + | T Float -> dual_number t ~primal:value ~tangent:(Some (zeros ())) + | T Bool -> dual_number t ~primal:value ~tangent:None) +;; + +let handle t ~f = + Fox_effect.handle ~f ~handle:(fun op -> + let result = + match Op.map op ~f:(lift t) with + | Unary (Neg, a) -> + dual_number + t + ~primal:Value.O.(-a.primal) + ~tangent:(Option.map a.tangent ~f:Value.neg) + | Unary (Sin, a) -> + dual_number + t + ~primal:(Value.sin a.primal) + ~tangent: + (Option.map a.tangent ~f:(fun tangent -> + Value.O.(Value.cos a.primal * tangent))) + | Unary (Cos, a) -> + dual_number + t + ~primal:(Value.cos a.primal) + ~tangent: + (Option.map a.tangent ~f:(fun tangent -> + Value.O.(-Value.sin a.primal * tangent))) + | Unary (Sqrt, a) -> + dual_number + t + ~primal:(Value.sqrt a.primal) + ~tangent: + (Option.map a.tangent ~f:(fun tangent -> + Value.div tangent (Value.scale (Value.sqrt a.primal) 2.))) + | Unary (Exp, a) -> + dual_number + t + ~primal:(Value.exp a.primal) + ~tangent: + (Option.map a.tangent ~f:(fun tangent -> + Value.O.(Value.exp a.primal * tangent))) + | Unary (Log, a) -> + dual_number + t + ~primal:(Value.log a.primal) + ~tangent:(Option.map a.tangent ~f:(fun tangent -> Value.div tangent a.primal)) + | Unary (Sigmoid, a) -> + dual_number + t + ~primal:(Value.sigmoid a.primal) + ~tangent: + (Option.map a.tangent ~f:(fun tangent -> + Value.O.( + Value.sigmoid a.primal + * ((Value.of_float 1. |> Value.broadcast ~dims:(Value.dims a.primal)) + - Value.sigmoid a.primal) + * tangent))) + | Binary (Add, a, b) -> + dual_number + t + ~primal:Value.O.(a.primal + b.primal) + ~tangent: + (Option.map2 a.tangent b.tangent ~f:(fun a_tangent b_tangent -> + Value.O.(a_tangent + b_tangent))) + | Binary (Sub, a, b) -> + dual_number + t + ~primal:Value.O.(a.primal - b.primal) + ~tangent: + (Option.map2 a.tangent b.tangent ~f:(fun a_tangent b_tangent -> + Value.O.(a_tangent - b_tangent))) + | Binary (Mul, a, b) -> + dual_number + t + ~primal:Value.O.(a.primal * b.primal) + ~tangent: + (Option.map2 a.tangent b.tangent ~f:(fun a_tangent b_tangent -> + Value.O.((a_tangent * b.primal) + (a.primal * b_tangent)))) + | Binary (Div, a, b) -> + dual_number + t + ~primal:Value.O.(a.primal / b.primal) + ~tangent: + (Option.map2 a.tangent b.tangent ~f:(fun a_tangent b_tangent -> + Value.O.( + ((a_tangent * b.primal) - (a.primal * b_tangent)) / (b.primal * b.primal)))) + | Binary (Eq, a, b) -> + dual_number t ~primal:Value.O.(a.primal = b.primal) ~tangent:None + | Binary (Gt, a, b) -> + dual_number t ~primal:Value.O.(a.primal > b.primal) ~tangent:None + | Binary (Lt, a, b) -> + dual_number t ~primal:Value.O.(a.primal < b.primal) ~tangent:None + | Matmul (a, b) -> + dual_number + t + ~primal:(Value.matmul a.primal b.primal) + ~tangent: + (Option.map2 a.tangent b.tangent ~f:(fun a_tangent b_tangent -> + Value.O.(Value.matmul a_tangent b.primal + Value.matmul a.primal b_tangent))) + | Transpose a -> + dual_number + t + ~primal:(Value.transpose a.primal) + ~tangent:(Option.map a.tangent ~f:Value.transpose) + | Sum { value; dims; keep_dims } -> + dual_number + t + ~primal:(Value.sum ~dims ~keep_dims value.primal) + ~tangent:(Option.map value.tangent ~f:(Value.sum ~dims ~keep_dims)) + | Broadcast { value; dims } -> + dual_number + t + ~primal:(Value.broadcast ~dims value.primal) + ~tangent:(Option.map value.tangent ~f:(Value.broadcast ~dims)) + | Reshape { value; dims } -> + dual_number + t + ~primal:(Value.reshape value.primal ~dims) + ~tangent:(Option.map value.tangent ~f:(Value.reshape ~dims)) + in + Dual_number.to_value result) +;; diff --git a/lib/core/jvp.mli b/lib/core/jvp.mli new file mode 100644 index 0000000..42ad95d --- /dev/null +++ b/lib/core/jvp.mli @@ -0,0 +1,31 @@ +open! Core + +module Dual_number : sig + (* XCR mtakeda: Does this type need to be exposed? Could we e.g. make it abstract with + creation functions / private? user: good call - made it abstract. The only consumer + (the [jvp] transform) reads [primal]/[tangent] and calls [to_value], so those + accessors are all we expose; the internal [id] tag (and the representation) are now + hidden. No external code constructs a [Dual_number] directly - they come from + [dual_number] / [lift] - so no creation function is needed here either. *) + + (** A primal value paired with its tangent (absent for non-differentiable values). + Carried through a computation as a [Value.t] - see [to_value]. *) + type t + + val primal : t -> Value.t + val tangent : t -> Value.t option + val to_value : t -> Value.t +end + +(** The forward-mode (JVP) tracer. Under [handle], each [Op] effect is interpreted on + [Dual_number]s, propagating tangents alongside primals. *) +type t + +val create : unit -> t +val dual_number : t -> primal:Value.t -> tangent:Value.t option -> Dual_number.t + +(** Views an arbitrary [Value.t] as a [Dual_number]: a value already tagged by this tracer + keeps its tangent, anything else gets a zero (or absent, for non-float) tangent. *) +val lift : t -> Value.t -> Dual_number.t + +val handle : t -> f:(unit -> 'a) -> 'a diff --git a/lib/core/op.ml b/lib/core/op.ml index 96d8ae0..4ac6183 100644 --- a/lib/core/op.ml +++ b/lib/core/op.ml @@ -9,7 +9,7 @@ module Unary = struct | Exp | Log | Sigmoid - [@@deriving sexp, enumerate] + [@@deriving sexp, enumerate, compare, hash] let to_string t = [%sexp_of: t] t |> Sexp.to_string |> String.lowercase end @@ -23,7 +23,7 @@ module Binary = struct | Eq | Gt | Lt - [@@deriving sexp, enumerate] + [@@deriving sexp, enumerate, compare, hash] let to_string t = [%sexp_of: t] t |> Sexp.to_string |> String.lowercase end @@ -40,13 +40,13 @@ type 'value t = } | Broadcast of { value : 'value - ; dims : int array + ; dims : int iarray } | Reshape of { value : 'value - ; dims : int array + ; dims : int iarray } -[@@deriving sexp_of, variants] +[@@deriving sexp_of, compare, hash, variants] let map t ~f = match t with @@ -123,12 +123,12 @@ let to_string t ~f = [%string "sum %{f value} dims=%{dims} keep_dims=%{keep_dims#Bool}"] | Broadcast { value; dims } -> let dims = - Array.to_list dims |> List.map ~f:Int.to_string |> String.concat ~sep:", " + Iarray.to_list dims |> List.map ~f:Int.to_string |> String.concat ~sep:", " in [%string "broadcast %{f value} dims=[%{dims}]"] | Reshape { value; dims } -> let dims = - Array.to_list dims |> List.map ~f:Int.to_string |> String.concat ~sep:", " + Iarray.to_list dims |> List.map ~f:Int.to_string |> String.concat ~sep:", " in [%string "reshape %{f value} dims=[%{dims}]"] ;; @@ -146,7 +146,7 @@ let infer_shape (t : Shape.t t) : Shape.t Or_error.t = , { dims = dims1; type_ = type1 } , { dims = dims2; type_ = type2 } ) -> let%map.Or_error () = - if [%equal: int array] dims1 dims2 + if [%equal: int iarray] dims1 dims2 then Ok () else Or_error.error_s [%message "infer_dims: dims mismatch" ~op:(t : Shape.t t)] and () = @@ -176,8 +176,8 @@ let infer_shape (t : Shape.t t) : Shape.t Or_error.t = in let%map.Or_error dims = match dims1, dims2 with - | [| n; m |], [| m' |] -> if m <> m' then dim_mismatch () else Ok [| n |] - | [| n; m |], [| m'; k |] -> if m <> m' then dim_mismatch () else Ok [| n; k |] + | [: n; m :], [: m' :] -> if m <> m' then dim_mismatch () else Ok [: n :] + | [: n; m :], [: m'; k :] -> if m <> m' then dim_mismatch () else Ok [: n; k :] | _ -> Or_error.error_s [%message @@ -187,7 +187,7 @@ let infer_shape (t : Shape.t t) : Shape.t Or_error.t = | Transpose { dims; type_ } -> let%map.Or_error dims = match dims with - | [| n; k |] -> Ok [| k; n |] + | [: n; k :] -> Ok [: k; n :] | _ -> Or_error.error_s [%message @@ -204,7 +204,7 @@ let infer_shape (t : Shape.t t) : Shape.t Or_error.t = in let%map.Or_error dims = match dims_to_sum with - | `All -> Ok (if keep_dims then Array.map dims ~f:(fun _ -> 1) else [||]) + | `All -> Ok (if keep_dims then Iarray.map dims ~f:(fun _ -> 1) else [::]) | `Just dims_to_sum -> let%bind.Or_error () = if Nonempty_list.to_list dims_to_sum |> List.contains_dup ~compare:Int.compare @@ -214,7 +214,7 @@ let infer_shape (t : Shape.t t) : Shape.t Or_error.t = "infer_dims: Sum: duplicate reduction dimension" ~op:(t : Shape.t t)] else Ok () in - let dims_length = Array.length dims in + let dims_length = Iarray.length dims in (match Nonempty_list.for_all dims_to_sum ~f:(fun dim -> dim < dims_length || dims_length + dim >= 0) @@ -232,58 +232,57 @@ let infer_shape (t : Shape.t t) : Shape.t Or_error.t = in if keep_dims then - Array.mapi dims ~f:(fun index dim -> + Iarray.mapi dims ~f:(fun index dim -> if Set.mem dims_to_sum index then 1 else dim) else - Array.filteri dims ~f:(fun index _dim -> not (Set.mem dims_to_sum index)) + Iarray.filteri dims ~f:(fun index _dim -> not (Set.mem dims_to_sum index)) in Ok dims) in { Shape.dims; type_ = T Float } | Broadcast { value = { dims = from_dims; type_ }; dims = to_dims } -> let%bind.Or_error () = - if Array.for_all to_dims ~f:Int.is_positive + if Iarray.for_all to_dims ~f:Int.is_positive then Ok () else Or_error.error_s [%message "infer_dims: dims must be positive" ~op:(t : Shape.t t)] in let%bind.Or_error () = - if Array.length to_dims >= Array.length from_dims + if Iarray.length to_dims >= Iarray.length from_dims then Ok () else Or_error.error_s [%message "infer_dims: can't broadcast to a larger rank" ~op:(t : Shape.t t)] in let%bind.Or_error () = - let dims_padding_length = Array.length to_dims - Array.length from_dims in + let dims_padding_length = Iarray.length to_dims - Iarray.length from_dims in let padded_from_dims = - Array.append (Array.create ~len:dims_padding_length 1) from_dims + Iarray.append (Iarray.init dims_padding_length ~f:(fun _ -> 1)) from_dims in - if - Array.zip_exn padded_from_dims to_dims - |> Array.for_all ~f:(fun (from, to_) -> to_ = from || from = 1) + if Iarray.zip_exn padded_from_dims to_dims + |> Iarray.for_all ~f:(fun (from, to_) -> to_ = from || from = 1) then Ok () else Or_error.error_s [%message "infer_dims: can't broadcast" ~op:(t : Shape.t t)] in Ok { Shape.dims = to_dims; type_ } | Reshape { value = { dims = from_dims; type_ }; dims = to_dims } -> - let from_dims_elements = Array.fold from_dims ~init:1 ~f:( * ) in + let from_dims_elements = Iarray.fold from_dims ~init:1 ~f:( * ) in let%map.Or_error dims = - match Array.count to_dims ~f:(fun dim -> dim = -1) with + match Iarray.count to_dims ~f:(fun dim -> dim = -1) with | 0 -> - let to_dims_elements = Array.fold to_dims ~init:1 ~f:( * ) in + let to_dims_elements = Iarray.fold to_dims ~init:1 ~f:( * ) in if from_dims_elements <> to_dims_elements then Or_error.error_s [%message "infer_dims: can't reshape" ~op:(t : Shape.t t)] else Ok to_dims | 1 -> let length_without_unknown_dim = - Array.filter to_dims ~f:(fun dim -> dim <> -1) |> Array.fold ~init:1 ~f:( * ) + Iarray.filter to_dims ~f:(fun dim -> dim <> -1) |> Iarray.fold ~init:1 ~f:( * ) in (match from_dims_elements % length_without_unknown_dim with | 0 -> let dims = - Array.map to_dims ~f:(fun dim -> + Iarray.map to_dims ~f:(fun dim -> if dim = -1 then from_dims_elements / length_without_unknown_dim else dim) in Ok dims @@ -307,8 +306,7 @@ module Make_operators (M : sig val eval : t op -> t val shape : t -> Shape.t end) : Operators_intf.S with type t := M.t = struct - let eval = - fun t -> + let eval t = let inferred_out_shape = map t ~f:M.shape |> infer_shape_exn in let out = M.eval t in [%test_result: Shape.t] @@ -355,7 +353,7 @@ module Make_operators (M : sig let dims value = M.shape value |> Shape.dims let scale value float = O.(value * broadcast (M.of_float float) ~dims:(dims value)) - let length value = Array.fold (dims value) ~init:1 ~f:( * ) + let length value = Iarray.fold (dims value) ~init:1 ~f:( * ) let mean ?dims:over_dims ?keep_dims value = let sum = sum ?dims:over_dims ?keep_dims value in diff --git a/lib/core/op.mli b/lib/core/op.mli index 293805f..8dab270 100644 --- a/lib/core/op.mli +++ b/lib/core/op.mli @@ -9,7 +9,7 @@ module Unary : sig | Exp | Log | Sigmoid - [@@deriving sexp, enumerate] + [@@deriving sexp, enumerate, compare, hash] end module Binary : sig @@ -21,7 +21,7 @@ module Binary : sig | Eq | Gt | Lt - [@@deriving sexp, enumerate] + [@@deriving sexp, enumerate, compare, hash] end type 'value t = @@ -36,13 +36,13 @@ type 'value t = } | Broadcast of { value : 'value - ; dims : int array + ; dims : int iarray } | Reshape of { value : 'value - ; dims : int array + ; dims : int iarray } -[@@deriving sexp_of] +[@@deriving sexp_of, compare, hash] val map : 'a t -> f:('a -> 'b) -> 'b t val to_list : 'a t -> 'a list diff --git a/lib/core/operators_intf.ml b/lib/core/operators_intf.ml index b27a64d..a5932ae 100644 --- a/lib/core/operators_intf.ml +++ b/lib/core/operators_intf.ml @@ -19,7 +19,7 @@ module type S = sig val lt : t -> t -> t val matmul : t -> t -> t val transpose : t -> t - val reshape : t -> dims:int array -> t + val reshape : t -> dims:int iarray -> t val sum : ?dims:[ `Just of int Nonempty_list.t | `All ] -> ?keep_dims:bool -> t -> t val mean : ?dims:[ `Just of int Nonempty_list.t | `All ] -> ?keep_dims:bool -> t -> t @@ -38,7 +38,7 @@ module type S = sig -> t val softmax : dim:int -> t -> t - val broadcast : t -> dims:int array -> t + val broadcast : t -> dims:int iarray -> t val scale : t -> float -> t module O : sig diff --git a/lib/core/partial_eval.ml b/lib/core/partial_eval.ml new file mode 100644 index 0000000..565311f --- /dev/null +++ b/lib/core/partial_eval.ml @@ -0,0 +1,116 @@ +open! Core + +module Partial_value = struct + type t = + | Known of Value.t + | Unknown of Expr.Var.t + [@@deriving sexp_of] + + let shape = function + | Known value -> Value.shape value + | Unknown var -> Expr.Var.shape var + ;; + + let to_atom t ~vars = + match t with + | Known value -> Expr.Atom.of_value value ~vars + | Unknown var -> Expr.Atom.Var var + ;; + + let type_id = Type_equal.Id.create ~name:"Partial_value" [%sexp_of: t] +end + +(* The partial-evaluation tracer. [Known] inputs are evaluated through; an op with any + [Unknown] operand is recorded as an equation and produces a fresh [Unknown]. *) +module Tracer = struct + type t = + { mutable equations : Value.t Expr.Eq.t list + ; mutable name_counter : int + ; mutable vars : Expr.Var.Set.t + } + + let create () = { equations = []; name_counter = 0; vars = Expr.Var.Set.empty } + + let fresh_var t ~shape = + let name = [%string "p_%{t.name_counter#Int}"] in + t.name_counter <- t.name_counter + 1; + let var : Expr.Var.t = { name; shape } in + t.vars <- Set.add t.vars var; + var + ;; + + let lift value = + Value.coerce value ~type_id:Partial_value.type_id + |> Option.value ~default:(Partial_value.Known value) + ;; + + let value_to_atom t value = Partial_value.to_atom value ~vars:t.vars + + let handle t ~f = + Fox_effect.handle ~f ~handle:(fun op -> + let result : Partial_value.t = + match Op.map op ~f:lift with + | Unary (kind, Known a) -> Known (Op.eval (module Value) (Unary (kind, a))) + | Binary (kind, Known a, Known b) -> + Known (Op.eval (module Value) (Binary (kind, a, b))) + | Matmul (Known a, Known b) -> Known (Value.matmul a b) + | Transpose (Known a) -> Known (Value.transpose a) + | Sum { value = Known a; dims; keep_dims } -> Known (Value.sum a ~dims ~keep_dims) + | Broadcast { value = Known a; dims } -> Known (Value.broadcast a ~dims) + | Reshape { value = Known a; dims } -> Known (Value.reshape a ~dims) + | ( Unary ((Neg | Sin | Cos | Sqrt | Exp | Log | Sigmoid), _) + | Binary ((Add | Sub | Mul | Div | Eq | Gt | Lt), _, _) + | Matmul _ | Transpose _ | Sum _ | Broadcast _ | Reshape _ ) as op -> + let shape = Op.map op ~f:Partial_value.shape |> Op.infer_shape_exn in + let binder = fresh_var t ~shape in + t.equations + <- { var = binder; op = Op.map op ~f:(value_to_atom t) } :: t.equations; + Unknown binder + in + Value.create + ~value:result + ~type_id:Partial_value.type_id + ~shape:(Partial_value.shape result)) + ;; + + let equations t = List.rev t.equations +end + +let partially_apply_expr_flat + (inputs : Partial_value.t list) + ~(f : Value.t list -> Value.t list * Value_tree.Def.t) + : Partial_value.t list * Value.t Expr.t + = + let partial = Tracer.create () in + let outputs, out_tree_def = + Tracer.handle partial ~f:(fun () -> + List.map inputs ~f:(fun input -> + Value.create + ~value:input + ~type_id:Partial_value.type_id + ~shape:(Partial_value.shape input)) + |> f) + in + let outputs = List.map outputs ~f:Tracer.lift in + let only_unknowns = + List.filter_map ~f:(function + | Partial_value.Known _ -> None + | Unknown var -> Some var) + in + ( outputs + , match + Expr.create + ~parameters:(only_unknowns inputs) + ~consts:Expr.Var.Map.empty + ~equations:(Tracer.equations partial) + ~return_vals: + (only_unknowns outputs + |> List.map ~f:(fun var -> Expr.Atom.Var var) + |> Nonempty_list.of_list_exn) + ~out_tree_def + with + | exception exn -> + raise_s + [%message "Failed to create expr" (exn : exn) (inputs : Partial_value.t list)] + | expr -> expr ) +;; diff --git a/lib/core/partial_eval.mli b/lib/core/partial_eval.mli new file mode 100644 index 0000000..37b2b45 --- /dev/null +++ b/lib/core/partial_eval.mli @@ -0,0 +1,20 @@ +open! Core + +module Partial_value : sig + (** A value during partial evaluation: either fully [Known], or an [Unknown] standing + for a traced variable. *) + type t = + | Known of Value.t + | Unknown of Expr.Var.t + [@@deriving sexp_of] + + val shape : t -> Shape.t +end + +(** Partially evaluates [f] over [inputs]: [Known] inputs are computed through while + [Unknown] inputs are traced, yielding the partial outputs and an [Expr.t] capturing + the traced (unknown-dependent) part of the computation. *) +val partially_apply_expr_flat + : Partial_value.t list + -> f:(Value.t list -> Value.t list * Value_tree.Def.t) + -> Partial_value.t list * Value.t Expr.t diff --git a/lib/core/shape.ml b/lib/core/shape.ml index f3d5717..660baec 100644 --- a/lib/core/shape.ml +++ b/lib/core/shape.ml @@ -1,12 +1,9 @@ open! Core -module T = struct - type t = - { dims : int array - ; type_ : Type.Packed.t - } - [@@deriving equal, compare, sexp, fields ~getters] -end +type t = + { dims : int iarray + ; type_ : Type.Packed.t + } +[@@deriving equal, compare, hash, sexp, fields ~getters] -include T -include Comparable.Make_plain (T) +include functor Comparable.Make_plain diff --git a/lib/core/shape.mli b/lib/core/shape.mli index 0c581ee..d260acc 100644 --- a/lib/core/shape.mli +++ b/lib/core/shape.mli @@ -1,9 +1,9 @@ open! Core type t = - { dims : int array + { dims : int iarray ; type_ : Type.Packed.t } -[@@deriving equal, compare, sexp, fields ~getters] +[@@deriving equal, compare, hash, sexp, fields ~getters] include Comparable.S_plain with type t := t diff --git a/lib/core/staging.ml b/lib/core/staging.ml new file mode 100644 index 0000000..63881f1 --- /dev/null +++ b/lib/core/staging.ml @@ -0,0 +1,55 @@ +open! Core + +type t = + { mutable equations : Value.t Expr.Eq.t list + ; mutable var_name_counter : int + ; mutable const_name_counter : int + ; mutable vars : Expr.Var.Set.t + ; mutable consts : Expr.Var.t Value.On_id.Map.t + } + +let create () = + { equations = [] + ; var_name_counter = 0 + ; const_name_counter = 0 + ; vars = Expr.Var.Set.empty + ; consts = Value.On_id.Map.empty + } +;; + +let fresh_var t ~shape = + let name = [%string "v_%{t.var_name_counter#Int}"] in + t.var_name_counter <- t.var_name_counter + 1; + let var : Expr.Var.t = { name; shape } in + t.vars <- Set.add t.vars var; + var +;; + +let intern_value t value : Value.t Expr.Atom.t = + match Expr.Atom.of_value ~vars:t.vars value with + | Var var -> Var var + | Value value -> + (match Map.find t.consts value with + | Some const -> Var const + | None -> + let name = [%string "c_%{t.const_name_counter#Int}"] in + t.const_name_counter <- t.const_name_counter + 1; + let const : Expr.Var.t = { name; shape = Value.shape value } in + t.consts <- Map.add_exn t.consts ~key:value ~data:const; + t.vars <- Set.add t.vars const; + Var const) +;; + +let handle t ~f = + Fox_effect.handle ~f ~handle:(fun op -> + let shape = Op.map op ~f:Value.shape |> Op.infer_shape_exn in + let binder = fresh_var t ~shape in + t.equations <- { var = binder; op = Op.map op ~f:(intern_value t) } :: t.equations; + Value.create ~value:binder ~type_id:Expr.Var.type_id ~shape) +;; + +let equations t = List.rev t.equations + +let consts_map t = + Map.to_alist t.consts |> List.map ~f:Tuple2.swap |> Expr.Var.Map.of_alist_exn +;; diff --git a/lib/core/staging.mli b/lib/core/staging.mli new file mode 100644 index 0000000..f072aff --- /dev/null +++ b/lib/core/staging.mli @@ -0,0 +1,21 @@ +open! Core + +(** The tracer behind [build_expr]. Under [handle], each [Op] effect is recorded as an + equation over fresh variables, and the constants it touches are hoisted (deduplicated + by identity) into named const-vars. *) +type t + +val create : unit -> t +val fresh_var : t -> shape:Shape.t -> Expr.Var.t + +(** Resolves a value to an atom, hoisting it into a fresh (or shared) const-var if it is a + constant rather than one of this tracer's variables. *) +val intern_value : t -> Value.t -> Value.t Expr.Atom.t + +val handle : t -> f:(unit -> 'a) -> 'a + +(** The recorded equations, in evaluation order. *) +val equations : t -> Value.t Expr.Eq.t list + +(** The hoisted constants, as a map from each const-var to its value. *) +val consts_map : t -> Value.t Expr.Var.Map.t diff --git a/lib/core/tensor.ml b/lib/core/tensor.ml index 25cf4e8..a2864e4 100644 --- a/lib/core/tensor.ml +++ b/lib/core/tensor.ml @@ -13,7 +13,9 @@ module Typed = struct } -> 'logical_type t - let dims (T { bigarray; mapping = _ }) = Bigarray.Genarray.dims bigarray + let dims (T { bigarray; mapping = _ }) = + Iarray.of_array (Bigarray.Genarray.dims bigarray) + ;; let type_ (type a) (T { bigarray = _; mapping } : a t) : a Type.t = match mapping with @@ -23,7 +25,7 @@ module Typed = struct let shape t : Shape.t = { dims = dims t; type_ = T (type_ t) } let num_dims (T { bigarray; mapping = _ }) = Bigarray.Genarray.num_dims bigarray - let length t = dims t |> Array.fold ~init:1 ~f:( * ) + let length t = dims t |> Iarray.fold ~init:1 ~f:( * ) let char_to_bool = function | '\000' -> false @@ -37,21 +39,23 @@ module Typed = struct ;; let item (type a) (T { bigarray; mapping } as t : a t) : a = - match dims t with - | [||] -> + match Iarray.length (dims t) with + | 0 -> (match mapping with | Float -> Bigarray.Genarray.get bigarray [||] | Bool -> Bigarray.Genarray.get bigarray [||] |> char_to_bool) - | dims -> raise_s [%message "Tensor.item: dims > 0" (dims : int array)] + | _ -> raise_s [%message "Tensor.item: dims > 0" ~dims:(dims t : int iarray)] ;; let get (type a) (T { bigarray; mapping } : a t) index : a = + let index = Iarray.to_array index in match mapping with | Float -> Bigarray.Genarray.get bigarray index | Bool -> Bigarray.Genarray.get bigarray index |> char_to_bool ;; let set (type a) (T { bigarray; mapping } : a t) index (value : a) = + let index = Iarray.to_array index in match mapping with | Float -> Bigarray.Genarray.set bigarray index value | Bool -> Bigarray.Genarray.set bigarray index (char_of_bool value) @@ -64,7 +68,10 @@ module Typed = struct ;; let left_slice (type a) (T { bigarray; mapping } : a t) ~indices : a t = - T { bigarray = Bigarray.Genarray.slice_left bigarray indices; mapping } + T + { bigarray = Bigarray.Genarray.slice_left bigarray (Iarray.to_array indices) + ; mapping + } ;; let sexp_of_a (type a) (t : a t) : a -> Sexp.t = @@ -76,18 +83,20 @@ module Typed = struct let rec sexp_of_t : 'a. ('a -> Sexp.t) -> 'a t -> Sexp.t = fun sexp_of_a t -> match dims t with - | [||] -> item t |> [%sexp_of: a] - | [| n |] -> List.init n ~f:(fun i -> get t [| i |]) |> [%sexp_of: a list] + | [::] -> item t |> [%sexp_of: a] + | [: n :] -> List.init n ~f:(fun i -> get t [: i :]) |> [%sexp_of: a list] | dims -> - let first_dim = dims.(0) in + let first_dim = dims.:(0) in List.init first_dim ~f:(fun i -> - let t = left_slice t ~indices:[| i |] in + let t = left_slice t ~indices:[: i :] in sexp_of_t sexp_of_a t) |> [%sexp_of: Sexp.t list] ;; let sexp_of_t sexp_of_a t = - if length t > 30 then [%sexp { dims : int array = dims t }] else sexp_of_t sexp_of_a t + if length t > 30 + then [%sexp { dims : int iarray = dims t }] + else sexp_of_t sexp_of_a t ;; include Type_equal.Id.Create1 (struct @@ -110,18 +119,21 @@ module Typed = struct } ;; - let init (type a) (type_ : a Type.t) ~dims ~(f : int array -> a) : a t = + let init (type a) (type_ : a Type.t) ~dims ~(f : int iarray -> a) : a t = + let dims = Iarray.to_array dims in match type_ with | Float -> T - { bigarray = Bigarray.Genarray.init Bigarray.float64 Bigarray.c_layout dims f + { bigarray = + Bigarray.Genarray.init Bigarray.float64 Bigarray.c_layout dims (fun index -> + f (Iarray.of_array index)) ; mapping = Float } | Bool -> T { bigarray = Bigarray.Genarray.init Bigarray.char Bigarray.c_layout dims (fun index -> - f index |> char_of_bool) + f (Iarray.of_array index) |> char_of_bool) ; mapping = Bool } ;; @@ -130,9 +142,10 @@ module Typed = struct T { bigarray = Bigarray.reshape bigarray dims; mapping } ;; - (* TODO: it's a bit sad that we duplicate this checking logic with - [Op.infer_dims]; is there a way to reduce the duplication here...? *) + (* TODO: it's a bit sad that we duplicate this checking logic with [Op.infer_dims]; is + there a way to reduce the duplication here...? *) let reshape t ~dims = + let dims = Iarray.to_array dims in match Array.count dims ~f:(fun dim -> dim = -1) with | 0 -> reshape' t ~dims | 1 -> @@ -159,11 +172,11 @@ module Typed = struct ;; (* TODO: write a ppx that allows writing [5t] or [5.5t] which expands to - [Tensor.of_float (Int.to_float 5)] and [Tensor.of_float 5.5]. See - janestreet/ppx_fixed_literal for prior art. *) + [Tensor.of_float (Int.to_float 5)] and [Tensor.of_float 5.5]. See + janestreet/ppx_fixed_literal for prior art. *) let of_lit (type a) (type_ : a Type.t) f = let t = create_uninitialized type_ [||] in - set t [||] f; + set t [::] f; t ;; @@ -175,7 +188,7 @@ module Typed = struct let of_list (type a) (type_ : a Type.t) (l : a list) : a t = let t = create_uninitialized type_ [| List.length l |] in - List.iteri l ~f:(fun i x -> set t [| i |] x); + List.iteri l ~f:(fun i x -> set t [: i :] x); t ;; @@ -189,7 +202,7 @@ module Typed = struct match List.map ~f:List.length l |> List.dedup_and_sort ~compare:Int.compare with | [ row_length ] -> let t = create_uninitialized type_ [| List.length l; row_length |] in - List.iteri l ~f:(fun i row -> List.iteri row ~f:(fun j x -> set t [| i; j |] x)); + List.iteri l ~f:(fun i row -> List.iteri row ~f:(fun j x -> set t [: i; j :] x)); t | row_lengths -> raise_s [%message "of_list2_exn: non-rectangular list" (row_lengths : int list)] @@ -202,7 +215,7 @@ module Typed = struct ;; let create (type a) (type_ : a Type.t) ~dims (value : a) : a t = - let t = create_uninitialized type_ dims in + let t = create_uninitialized type_ (Iarray.to_array dims) in fill t value; t ;; @@ -213,7 +226,7 @@ module Typed = struct let arange n = let t = create_uninitialized Float [| n |] in for i = 0 to n - 1 do - set t [| i |] (Int.to_float i) + set t [: i :] (Int.to_float i) done; t ;; @@ -221,9 +234,9 @@ module Typed = struct let%expect_test "arange" = arange 12 |> [%sexp_of: float t] |> print_s; [%expect {| (0 1 2 3 4 5 6 7 8 9 10 11) |}]; - arange 12 |> reshape ~dims:[| 6; 2 |] |> [%sexp_of: float t] |> print_s; + arange 12 |> reshape ~dims:[: 6; 2 :] |> [%sexp_of: float t] |> print_s; [%expect {| ((0 1) (2 3) (4 5) (6 7) (8 9) (10 11)) |}]; - arange 12 |> reshape ~dims:[| 3; 4 |] |> [%sexp_of: float t] |> print_s; + arange 12 |> reshape ~dims:[: 3; 4 :] |> [%sexp_of: float t] |> print_s; [%expect {| ((0 1 2 3) (4 5 6 7) (8 9 10 11)) |}] ;; @@ -233,10 +246,10 @@ module Typed = struct let map2 (type a) (type_ : a Type.t) t1 t2 ~f = let dims1 = dims t1 in let dims2 = dims t2 in - if not ([%compare.equal: int array] dims1 dims2) + if not ([%compare.equal: int iarray] dims1 dims2) then raise_s - [%message "Tensor.map2: dims mismatch" (dims1 : int array) (dims2 : int array)]; + [%message "Tensor.map2: dims mismatch" (dims1 : int iarray) (dims2 : int iarray)]; init type_ ~dims:dims1 ~f:(fun index -> f (get t1 index) (get t2 index)) ;; @@ -253,7 +266,7 @@ module Typed = struct let iter t ~f = iteri t ~f:(fun _index value -> f value) let allclose (type a) ?(equal_nan = false) (t1 : a t) (t2 : a t) = - [%equal: int array] (dims t1) (dims t2) + [%equal: int iarray] (dims t1) (dims t2) && let is_equal = ref true in iteri t1 ~f:(fun index value1 -> @@ -270,7 +283,7 @@ module Typed = struct ;; let sum_single_axis t ~axis ~keep_dim = - let dims = dims t in + let dims = Iarray.to_array (dims t) in let dims_length = Array.length dims in if axis < 0 || axis >= dims_length then raise_s [%message "sum_single_axis: axis out of bounds" (axis : int)]; @@ -279,10 +292,12 @@ module Typed = struct let result_dims = Array.concat [ dims_left; (if keep_dim then [| 1 |] else [||]); dims_right ] in - init Float ~dims:result_dims ~f:(fun index -> + init Float ~dims:(Iarray.of_array result_dims) ~f:(fun index -> + (* [Iarray.to_array] returns a fresh array we can mutate as scratch for [get]. *) + let index = Iarray.to_array index in let index = if keep_dim - then Array.copy index + then index else Array.init dims_length ~f:(fun i -> match Ordering.of_int (Int.compare i axis) with @@ -293,7 +308,7 @@ module Typed = struct let acc = ref 0. in for i = 0 to dims.(axis) - 1 do index.(axis) <- i; - acc := !acc +. get t index + acc := !acc +. get t (Iarray.of_array index) done; !acc) ;; @@ -414,47 +429,48 @@ let eval_op (op : t Op.t) = | Float, Float -> (* TODO: support more than just 2D tensors for matmuls and transposes *) (match Typed.dims t1, Typed.dims t2 with - | [| n; m |], [| m' |] -> + | [: n; m :], [: m' :] -> [%test_eq: int] m m'; let t = Typed.create_uninitialized Float [| n |] in for i = 0 to n - 1 do let acc = ref 0. in for l = 0 to m - 1 do - acc := !acc +. (Typed.get t1 [| i; l |] *. Typed.get t2 [| l |]) + acc := !acc +. (Typed.get t1 [: i; l :] *. Typed.get t2 [: l :]) done; - Typed.set t [| i |] !acc + Typed.set t [: i :] !acc done; T t - | [| n; m |], [| m'; k |] -> + | [: n; m :], [: m'; k :] -> [%test_eq: int] m m'; let t = Typed.create_uninitialized Float [| n; k |] in for i = 0 to n - 1 do for j = 0 to k - 1 do let acc = ref 0. in for l = 0 to m - 1 do - acc := !acc +. (Typed.get t1 [| i; l |] *. Typed.get t2 [| l; j |]) + acc := !acc +. (Typed.get t1 [: i; l :] *. Typed.get t2 [: l; j :]) done; - Typed.set t [| i; j |] !acc + Typed.set t [: i; j :] !acc done done; T t | t1_dims, t2_dims -> raise_s [%message - "matmul: unsupported dimensions" (t1_dims : int array) (t2_dims : int array)])) + "matmul: unsupported dimensions" + (t1_dims : int iarray) + (t2_dims : int iarray)])) | Transpose (T t) -> (match Typed.dims t with - | [| n; m |] -> + | [: n; m :] -> T - (Typed.init (Typed.type_ t) ~dims:[| m; n |] ~f:(fun index -> - Typed.get t [| index.(1); index.(0) |])) - | dims -> raise_s [%message "transpose: unsupported dimensions" (dims : int array)]) + (Typed.init (Typed.type_ t) ~dims:[: m; n :] ~f:(fun index -> + Typed.get t [: index.:(1); index.:(0) :])) + | dims -> raise_s [%message "transpose: unsupported dimensions" (dims : int iarray)]) | Sum { value = T t; dims = dims_to_sum; keep_dims } -> (match Typed.type_ t with | Bool -> raise_s [%message "eval_op: bool tensors not supported" (op : t Op.t)] | Float -> - let dims = Typed.dims t in - let dims_length = Array.length dims in + let dims_length = Iarray.length (Typed.dims t) in let dims_to_sum = (match dims_to_sum with | `Just dims_to_sum -> @@ -468,16 +484,17 @@ let eval_op (op : t Op.t) = (List.fold dims_to_sum ~init:t ~f:(fun t axis -> Typed.sum_single_axis t ~axis ~keep_dim:keep_dims))) | Broadcast { value = T t; dims = to_dims } -> - let from_dims = Typed.dims t in - let dims_padding_length = Array.length to_dims - Array.length from_dims in + let from_dims = Iarray.to_array (Typed.dims t) in + let dims_padding_length = Iarray.length to_dims - Array.length from_dims in T (Typed.init (Typed.type_ t) ~dims:to_dims ~f:(fun index -> + let index = Iarray.to_array index in let from_index = Array.subo index ~pos:dims_padding_length ~len:(Array.length from_dims) |> Array.map2_exn from_dims ~f:(fun from_dim index_dim -> if from_dim = 1 then 0 else index_dim) in - Typed.get t from_index)) + Typed.get t (Iarray.of_array from_index))) | Reshape { value = T t; dims = to_dims } -> T (Typed.reshape t ~dims:to_dims) ;; @@ -540,12 +557,12 @@ let%expect_test "std" = let%expect_test "broadcast" = let broadcast_and_print t ~dims:dims' = let t = broadcast t ~dims:dims' in - print_s [%message "" (t : t) ~dims:(dims t : int array)] + print_s [%message "" (t : t) ~dims:(dims t : int iarray)] in let t = of_list2_exn Float [ [ 1.; 2. ]; [ 3.; 4. ] ] in - broadcast_and_print t ~dims:[| 1; 2; 2 |]; + broadcast_and_print t ~dims:[: 1; 2; 2 :]; [%expect {| ((t (((1 2) (3 4)))) (dims (1 2 2))) |}]; - broadcast_and_print t ~dims:[| 2; 2; 2 |]; + broadcast_and_print t ~dims:[: 2; 2; 2 :]; [%expect {| ((t (((1 2) (3 4)) ((1 2) (3 4)))) (dims (2 2 2))) |}] ;; @@ -575,29 +592,29 @@ let normal ?(mean = 0.) ?(std = 1.) ~dims ~rng () = let%expect_test "normal" = let rng = Splittable_random.of_int 0 in - normal ~dims:[| 2; 2 |] ~rng () |> [%sexp_of: t] |> print_s; + normal ~dims:[: 2; 2 :] ~rng () |> [%sexp_of: t] |> print_s; [%expect {| ((0.39995642633665462 1.1602368073797789) (1.1461698444484156 -0.27508260258159217)) |}]; - let t = normal ~dims:[| 10000 |] ~rng () in + let t = normal ~dims:[: 10000 :] ~rng () in let mean = mean t |> item_exn Float in let std = std t |> item_exn Float in print_s [%message "" (mean : float) (std : float)]; - [%expect {| ((mean 0.0026463860836857677) (std 0.98795664491502377)) |}] + [%expect {| ((mean 0.0026463860836857694) (std 0.98795664491502377)) |}] ;; module With_shape = struct type nonrec t = t - let sexp_of_t t = [%sexp { dims : int array = dims t; tensor : t = t }] + let sexp_of_t t = [%sexp { dims : int iarray = dims t; tensor : t = t }] end module Just_shape = struct type nonrec t = t - let sexp_of_t t = [%sexp_of: int array] (dims t) + let sexp_of_t t = [%sexp_of: int iarray] (dims t) end module Private = struct diff --git a/lib/core/tensor.mli b/lib/core/tensor.mli index 0733594..780def9 100644 --- a/lib/core/tensor.mli +++ b/lib/core/tensor.mli @@ -7,33 +7,33 @@ module Typed : sig val type_ : 'a t -> 'a Type.t val type_equal_id : 'a Type_equal.Id.t -> 'a t Type_equal.Id.t - val dims : 'a t -> int array + val dims : 'a t -> int iarray val shape : 'a t -> Shape.t val num_dims : 'a t -> int val length : 'a t -> int val item : 'a t -> 'a - val get : 'a t -> int array -> 'a - val set : 'a t -> int array -> 'a -> unit + val get : 'a t -> int iarray -> 'a + val set : 'a t -> int iarray -> 'a -> unit val fill : 'a t -> 'a -> unit - val left_slice : 'a t -> indices:int array -> 'a t + val left_slice : 'a t -> indices:int iarray -> 'a t val sub_left : 'a t -> pos:int -> len:int -> 'a t val of_lit : 'a Type.t -> 'a -> 'a t val of_list : 'a Type.t -> 'a list -> 'a t - (** [of_list2_exn l] creates a tensor from a list of rows. Raises if a - non-rectangular list of lists are provided. *) + (** [of_list2_exn l] creates a tensor from a list of rows. Raises if a non-rectangular + list of lists are provided. *) val of_list2_exn : 'a Type.t -> 'a list list -> 'a t - val create : 'a Type.t -> dims:int array -> 'a -> 'a t - val init : 'a Type.t -> dims:int array -> f:(int array -> 'a) -> 'a t - val zeros : dims:int array -> float t - val ones : dims:int array -> float t + val create : 'a Type.t -> dims:int iarray -> 'a -> 'a t + val init : 'a Type.t -> dims:int iarray -> f:(int iarray -> 'a) -> 'a t + val zeros : dims:int iarray -> float t + val ones : dims:int iarray -> float t val arange : int -> float t val map : 'b Type.t -> 'a t -> f:('a -> 'b) -> 'b t - val mapi : 'b Type.t -> 'a t -> f:(int array -> 'a -> 'b) -> 'b t + val mapi : 'b Type.t -> 'a t -> f:(int iarray -> 'a -> 'b) -> 'b t val map2 : 'c Type.t -> 'a t -> 'b t -> f:('a -> 'b -> 'c) -> 'c t val iter : 'a t -> f:('a -> unit) -> unit - val iteri : 'a t -> f:(int array -> 'a -> unit) -> unit + val iteri : 'a t -> f:(int iarray -> 'a -> unit) -> unit val allclose : ?equal_nan:bool -> 'a t -> 'a t -> bool end @@ -41,20 +41,20 @@ type t = T : 'a Typed.t -> t [@@deriving sexp_of] val of_typed : 'a Typed.t -> t val to_typed_exn : 'a Type.t -> t -> 'a Typed.t -val dims : t -> int array +val dims : t -> int iarray val type_ : t -> Type.Packed.t val shape : t -> Shape.t val num_dims : t -> int val length : t -> int val item_exn : 'a Type.t -> t -> 'a -val get_exn : 'a Type.t -> t -> int array -> 'a -val left_slice : t -> indices:int array -> t +val get_exn : 'a Type.t -> t -> int iarray -> 'a +val left_slice : t -> indices:int iarray -> t val sub_left : t -> pos:int -> len:int -> t -val init : 'a Type.t -> dims:int array -> f:(int array -> 'a) -> t +val init : 'a Type.t -> dims:int iarray -> f:(int iarray -> 'a) -> t val of_list : 'a Type.t -> 'a list -> t val of_list2_exn : 'a Type.t -> 'a list list -> t val of_lit : 'a Type.t -> 'a -> t -val zeros : dims:int array -> t +val zeros : dims:int iarray -> t val allclose : ?equal_nan:bool -> t -> t -> bool include Operators_intf.S with type t := t @@ -62,7 +62,7 @@ include Operators_intf.S with type t := t val normal : ?mean:float -> ?std:float - -> dims:int array + -> dims:int iarray -> rng:Splittable_random.t -> unit -> t diff --git a/lib/core/treeable.ml b/lib/core/treeable.ml index 973a57d..ce7b01f 100644 --- a/lib/core/treeable.ml +++ b/lib/core/treeable.ml @@ -25,7 +25,7 @@ module Of_typed_fields (T : Of_typed_fields_arg) : S with type t := T.t = struct let t_of_tree (tree : Value_tree.t) = Typed_field.create { f = - (fun (type a) (field : a Typed_field.t) -> + (fun (type a) (local_ (field : a Typed_field.t)) -> let name = Typed_field.name field in let _, (module T) = field_treeable field in T.t_of_tree (Value_tree.get_exn tree name)) diff --git a/lib/core/treeable_intf.ml b/lib/core/treeable_intf.ml index 8fe65fb..df9498f 100644 --- a/lib/core/treeable_intf.ml +++ b/lib/core/treeable_intf.ml @@ -21,7 +21,7 @@ module type Of_typed_fields_arg = sig module Typed_field : Typed_fields_lib.S with type derived_on = t val field_treeable - : 'a Typed_field.t + : 'a Typed_field.t @ local -> ('a -> Value_tree.t) * (module S with type t = 'a) end diff --git a/lib/core/type.ml b/lib/core/type.ml index fcadedb..e69c9ae 100644 --- a/lib/core/type.ml +++ b/lib/core/type.ml @@ -32,6 +32,8 @@ module Packed = struct let compare (T t1) (T t2) = compare (rank t1) (rank t2) let equal = [%compare.equal: t] + let hash_fold_t state (T t) = hash_fold_int state (rank t) + let hash = Hash.of_fold hash_fold_t let sexp_of_t (T t) = [%sexp_of: _ typed] t let t_of_sexp sexp = diff --git a/lib/core/type.mli b/lib/core/type.mli index 5e5a155..84bb87c 100644 --- a/lib/core/type.mli +++ b/lib/core/type.mli @@ -9,5 +9,5 @@ val type_equal_id : 'a t -> 'a Type_equal.Id.t module Packed : sig type 'a typed := 'a t - type t = T : 'a typed -> t [@@deriving enumerate, equal, compare, sexp] + type t = T : 'a typed -> t [@@deriving enumerate, equal, compare, hash, sexp] end diff --git a/lib/core/value.ml b/lib/core/value.ml index 6b09f30..307984e 100644 --- a/lib/core/value.ml +++ b/lib/core/value.ml @@ -15,7 +15,7 @@ let of_typed_tensor (type a) (tensor : a Tensor.Typed.t) = let type_id = Tensor.Typed.type_ tensor |> Type.type_equal_id |> Tensor.Typed.type_equal_id in - T { value = tensor; type_id; shape = Tensor.Typed.shape tensor } + create ~value:tensor ~type_id ~shape:(Tensor.Typed.shape tensor) ;; let of_tensor (Tensor.T tensor) = of_typed_tensor tensor @@ -32,6 +32,7 @@ let to_tensor_exn t = let of_float x = of_typed_tensor (Tensor.Typed.of_lit Float x) let to_float_exn t : float = to_typed_tensor_exn Float t |> Tensor.Typed.item +let to_string t = [%sexp_of: t] t |> Sexp.to_string include Op.Make_operators (struct type nonrec t = t [@@deriving sexp_of] diff --git a/lib/core/value.mli b/lib/core/value.mli index 6418ed7..28c5661 100644 --- a/lib/core/value.mli +++ b/lib/core/value.mli @@ -1,26 +1,29 @@ open! Core -type t = Value0.t = +type t = Value0.t = private | T : { value : 'a ; type_id : 'a Type_equal.Id.t ; shape : Shape.t + ; id : Id.t } -> t [@@deriving sexp_of] -val dims : t -> int array +val create : value:'a -> type_id:'a Type_equal.Id.t -> shape:Shape.t -> t +val dims : t -> int iarray val type_ : t -> Type.Packed.t val shape : t -> Shape.t val coerce : t -> type_id:'a Type_equal.Id.t -> 'a option val coerce_exn : t -> type_id:'a Type_equal.Id.t -> 'a -val tree_def : dims:int array -> Value_tree.Def.t +val tree_def : dims:int iarray -> Value_tree.Def.t val of_typed_tensor : 'a Tensor.Typed.t -> t val of_tensor : Tensor.t -> t val to_typed_tensor_exn : 'a Type.t -> t -> 'a Tensor.Typed.t val to_tensor_exn : t -> Tensor.t val of_float : float -> t val to_float_exn : t -> float +val to_string : t -> string include Treeable.S with type t := t include Operators_intf.S with type t := t @@ -28,5 +31,7 @@ include Operators_intf.S with type t := t module Tuple2 : sig include Treeable.S with type t = t * t - val tree_def : dims1:int array -> dims2:int array -> Value_tree.Def.t + val tree_def : dims1:int iarray -> dims2:int iarray -> Value_tree.Def.t end + +module On_id : Comparable.S_plain with type t := t diff --git a/lib/core/value0.ml b/lib/core/value0.ml index e9790af..9d69d6d 100644 --- a/lib/core/value0.ml +++ b/lib/core/value0.ml @@ -1,23 +1,23 @@ open! Core -(* TODO: switch to use iarrays once they land: - https://github.com/ocaml/ocaml/pull/13097 *) type t = | T : { value : 'a ; type_id : 'a Type_equal.Id.t ; shape : Shape.t + ; id : Id.t } -> t -let dims (T { value = _; type_id = _; shape = { dims; type_ = _ } }) = dims -let type_ (T { value = _; type_id = _; shape = { dims = _; type_ } }) = type_ -let shape (T { value = _; type_id = _; shape }) = shape +let create ~value ~type_id ~shape = T { value; type_id; shape; id = Id.create () } +let dims (T { value = _; type_id = _; shape = { dims; type_ = _ }; id = _ }) = dims +let type_ (T { value = _; type_id = _; shape = { dims = _; type_ }; id = _ }) = type_ +let shape (T { value = _; type_id = _; shape; id = _ }) = shape let coerce - (type a) - (T { value; type_id = type_id'; shape = _ }) - ~(type_id : a Type_equal.Id.t) + (type a) + (T { value; type_id = type_id'; shape = _; id = _ }) + ~(type_id : a Type_equal.Id.t) : a option = match Type_equal.Id.same_witness type_id type_id' with @@ -26,24 +26,34 @@ let coerce ;; let coerce_exn - (type a) - (T { value; type_id = type_id'; shape = _ }) - ~(type_id : a Type_equal.Id.t) + (type a) + (T { value; type_id = type_id'; shape = _; id = _ }) + ~(type_id : a Type_equal.Id.t) : a = let T = Type_equal.Id.same_witness_exn type_id type_id' in value ;; -let sexp_of_t (T { value; type_id; shape = { dims; type_ } }) = +let sexp_of_t (T { value; type_id; shape = { dims; type_ }; id = _ }) = let x = Type_equal.Id.to_sexp type_id value in match dims with - | [||] -> + | [::] -> [%message (Type_equal.Id.name type_id) ~_:(x : Sexp.t) ~_:(type_ : Type.Packed.t)] - | dims -> + | _ -> [%message (Type_equal.Id.name type_id) ~_:(x : Sexp.t) - (dims : int array) - ~(type_ : Type.Packed.t)] + (dims : int iarray) + ~type_:(type_ : Type.Packed.t)] ;; + +(* Values carry a unique [id] so that an identical constant reused across a traced + computation can be deduplicated by identity. *) +module On_id = struct + type nonrec t = t [@@deriving sexp_of] + + let compare t1 t2 = Comparable.lift [%compare: Id.t] t1 t2 ~f:(fun (T { id; _ }) -> id) + + include functor Comparable.Make_plain +end diff --git a/lib/core/value0.mli b/lib/core/value0.mli index 4798206..6d51560 100644 --- a/lib/core/value0.mli +++ b/lib/core/value0.mli @@ -1,16 +1,21 @@ open! Core -type t = +type t = private | T : { value : 'a ; type_id : 'a Type_equal.Id.t ; shape : Shape.t + ; id : Id.t } -> t [@@deriving sexp_of] -val dims : t -> int array +val create : value:'a -> type_id:'a Type_equal.Id.t -> shape:Shape.t -> t +val dims : t -> int iarray val type_ : t -> Type.Packed.t val shape : t -> Shape.t val coerce : t -> type_id:'a Type_equal.Id.t -> 'a option val coerce_exn : t -> type_id:'a Type_equal.Id.t -> 'a + +(** Compares values by their unique [id]. *) +module On_id : Comparable.S_plain with type t := t diff --git a/lib/core/value_tree.ml b/lib/core/value_tree.ml index 2575dfb..6f0d7e0 100644 --- a/lib/core/value_tree.ml +++ b/lib/core/value_tree.ml @@ -4,7 +4,7 @@ module General = struct type 'value t = | Leaf of 'value | Node of 'value t Map.M(String).t - [@@deriving sexp, compare, variants, quickcheck] + [@@deriving sexp, compare, hash, variants, quickcheck] let rec length : _ t -> int = function | Leaf _ -> 1 @@ -43,7 +43,9 @@ let rec flatten : 'value General.t -> 'value list = function ;; module Def = struct - type t = int array General.t [@@deriving sexp, compare, quickcheck] + type t = int iarray General.t [@@deriving sexp, compare, hash] + + include functor Hashable.Make let leaf ~dims = General.leaf dims let node = General.node @@ -59,7 +61,7 @@ let rec unflatten' values ~(def : Def.t) ~sexp_of_value : _ General.t = | Leaf dims -> (match values with | [ value ] -> - [%test_result: int array] (Value0.dims value) ~expect:dims; + [%test_result: int iarray] (Value0.dims value) ~expect:dims; Leaf value | _ -> raise_s [%message "Expected singleton leaf value" (values : value list)]) | Node children -> diff --git a/lib/core/value_tree.mli b/lib/core/value_tree.mli index 817c075..c53c4f5 100644 --- a/lib/core/value_tree.mli +++ b/lib/core/value_tree.mli @@ -10,11 +10,13 @@ val to_value_exn : t -> Value0.t module Def : sig type t [@@deriving sexp_of, compare] + include Hashable.S with type t := t + (* TODO: change [dims] to [shape] *) - val leaf : dims:int array -> t + val leaf : dims:int iarray -> t val node : t String.Map.t -> t val length : t -> int - val flatten : t -> int array list + val flatten : t -> int iarray list end val to_def : t -> Def.t diff --git a/lib/jit/fox_jit.ml b/lib/jit/fox_jit.ml index e3f1da2..8fac634 100644 --- a/lib/jit/fox_jit.ml +++ b/lib/jit/fox_jit.ml @@ -18,12 +18,17 @@ let tensor_to_xla_literal tensor = ;; let xla_subcomp - ({ parameters; equations; return_vals; out_tree_def = _ } : Expr.t) - arguments - ~builder + ({ parameters; consts = _; equations; return_vals; out_tree_def = _ } : Value.t Expr.t) + arguments + ~const_ops + ~builder = - let env = List.zip_exn parameters arguments |> Expr.Var.Map.of_alist_exn in - let read_atom (atom : Expr.Atom.t) ~env = + let env = + List.zip_exn parameters arguments + |> Expr.Var.Map.of_alist_exn + |> Map.merge_disjoint_exn const_ops + in + let read_atom (atom : Value.t Expr.Atom.t) ~env = match atom with | Var var -> Map.find_exn env var | Value value -> @@ -48,7 +53,7 @@ let xla_subcomp | Sigmoid -> fun x -> let one = - Tensor.Typed.ones ~dims:(Xla.Op.dims x) + Tensor.Typed.ones ~dims:(Iarray.of_array (Xla.Op.dims x)) |> tensor_to_xla_literal |> Xla.Op.constant ~builder in @@ -76,7 +81,7 @@ let xla_subcomp (* TODO: support arbitrary dimensions *) Xla.Op.transpose a ~dim_indexes:[| 1; 0 |] | Sum { value = value, in_shape; dims; keep_dims } -> - let in_dims = Shape.dims in_shape in + let in_dims = Iarray.to_array (Shape.dims in_shape) in let dims = match dims with | `All -> Array.init (Array.length in_dims) ~f:Fn.id @@ -84,13 +89,15 @@ let xla_subcomp in Xla.Op.reduce_sum value ~dims ~keep_dims | Broadcast { value = value, in_shape; dims = out_dims } -> - let in_dims = Shape.dims in_shape in + let in_dims = Iarray.to_array (Shape.dims in_shape) in + let out_dims = Iarray.to_array out_dims in let padding_length = Array.length out_dims - Array.length in_dims in Xla.Op.broadcast_in_dim value ~out_dims ~broadcast_dims:(Array.mapi in_dims ~f:(fun i _ -> padding_length + i)) - | Reshape { value = value, _; dims } -> Xla.Op.reshape value ~dims + | Reshape { value = value, _; dims } -> + Xla.Op.reshape value ~dims:(Iarray.to_array dims) in let shape = Op.map op ~f:snd |> Op.infer_shape_exn in Map.add_exn env ~key:var ~data:(xla_op, shape)) @@ -105,7 +112,7 @@ let xla_subcomp let xla_builder = lazy (Xla.Builder.create ~name:"xla_call") -let xla_callable ?(print_hlo = false) (expr : Expr.t) = +let xla_callable ?(print_hlo = false) (expr : Value.t Expr.t) = let xla_builder = Lazy.force xla_builder in let xla_params = List.mapi expr.parameters ~f:(fun i { name; shape = { dims; type_ } as shape } -> @@ -117,11 +124,31 @@ let xla_callable ?(print_hlo = false) (expr : Expr.t) = (match type_ with | T Float -> F64 | T Bool -> Pred) - ~dims + ~dims:(Iarray.to_array dims) ~builder:xla_builder , shape )) in - let out, out_shapes = xla_subcomp expr xla_params ~builder:xla_builder in + (* Hoisted constants become XLA parameters (after the real parameters), supplied per + call at execution time rather than baked in, so one compiled executable serves any + constant values of the same structure. Their shapes come from the const vars, so + compilation depends only on the structure, not the values. *) + let num_params = List.length expr.parameters in + let const_ops = + Map.to_alist expr.consts + |> List.mapi ~f:(fun i (var, _value) -> + let shape = Expr.Var.shape var in + let op = + Xla.Op.parameter + (Expr.Var.name var) + ~id:(num_params + i) + ~ty:F64 + ~dims:(Iarray.to_array (Shape.dims shape)) + ~builder:xla_builder + in + var, (op, shape)) + |> Expr.Var.Map.of_alist_exn + in + let out, out_shapes = xla_subcomp expr xla_params ~const_ops ~builder:xla_builder in let xla_client = Xla.Client.cpu () in let xla_device = Xla.Client.addressable_devices xla_client |> List.hd_exn in let computation = Xla.Computation.build ~root:out in @@ -132,9 +159,9 @@ let xla_callable ?(print_hlo = false) (expr : Expr.t) = |> String.strip |> print_endline; let xla_exe = Xla.Executable.compile xla_client computation in - Staged.stage (fun inputs -> + Staged.stage (fun inputs const_values -> let inputs = - List.map inputs ~f:(fun tensor -> + List.map (inputs @ const_values) ~f:(fun tensor -> tensor_to_xla_literal tensor |> Xla.Buffer.of_host_literal ~device:xla_device) |> List.to_array in @@ -148,61 +175,86 @@ let xla_callable ?(print_hlo = false) (expr : Expr.t) = |> Nonempty_list.map ~f:(fun (literal, shape) -> tensor_of_xla_literal literal ~shape)) ;; -(* TODO: One crucial difference between the implementation here and autodidax is - that the compilation is not cached. One approach would be to have - [jit] take an [in_tree_def] argument, and return a staged function of type - [in_ -> out] (with some validation that the tree def matches up, or - alternatively recompile if it doesn't?). -*) +module Structure = struct + (* The value-free structure of a traced program - the compilation cache key. *) + type t = unit Expr.t [@@deriving compare, hash, sexp_of] +end + +(* [jit] returns a reusable function that traces [f] on every call and reuses a compiled + executable cached by the program's [Structure.t] (the traced [expr] with all values + erased). Because constants are fed as runtime parameters, one executable serves every + call that shares a structure, and each call supplies its own constants - so even a + closure whose captured constants change reuses the executable and still gets correct + results. + + This is the autodidax model (re-trace every call, cache only the compilation), rather + than keying on input shapes to skip re-tracing (real JAX). The structure key stays + sound even if value-dependent control flow is ever added, since distinct programs get + distinct keys. *) let jit - (type in_ out) - (module In : Treeable_intf.S with type t = in_) - (module Out : Treeable_intf.S with type t = out) - ?print_hlo - ~f - (input : in_) - : out + (type in_ out) + (module In : Treeable_intf.S with type t = in_) + (module Out : Treeable_intf.S with type t = out) + ?print_hlo + ~f + () + : (in_ -> out) Staged.t = - let input_tree = In.tree_of_t input in - let input_tree_def = Value_tree.to_def input_tree in - let flattened_input_tensors = - Value_tree.flatten input_tree |> List.map ~f:(Value.to_typed_tensor_exn Float) - in - let expr = build_expr (module In) (module Out) ~f ~in_tree_def:input_tree_def in - let xla_callable = xla_callable ?print_hlo expr |> Staged.unstage in - let output = - xla_callable flattened_input_tensors - |> Nonempty_list.to_list - |> List.map ~f:Value.of_tensor - in - Value_tree.unflatten output ~def:expr.out_tree_def |> Out.t_of_tree + let cache = Hashtbl.create (module Structure) in + Staged.stage (fun input -> + let input_tree = In.tree_of_t input in + let input_tree_def = Value_tree.to_def input_tree in + let flattened_input_tensors = + Value_tree.flatten input_tree |> List.map ~f:(Value.to_typed_tensor_exn Float) + in + let expr = build_expr (module In) (module Out) ~f ~in_tree_def:input_tree_def in + let callable = + Hashtbl.find_or_add + cache + (Expr.map expr ~f:(fun _ -> ())) + ~default:(fun () -> xla_callable ?print_hlo expr |> Staged.unstage) + in + let const_values = + Map.data expr.consts |> List.map ~f:(Value.to_typed_tensor_exn Float) + in + let output = + callable flattened_input_tensors const_values + |> Nonempty_list.to_list + |> List.map ~f:Value.of_tensor + in + Value_tree.unflatten output ~def:expr.out_tree_def |> Out.t_of_tree) ;; -let jit' ?print_hlo ~f x = jit (module Value) (module Value) ?print_hlo ~f x +let jit' ?print_hlo ~f () = jit (module Value) (module Value) ?print_hlo ~f () +let foo x = Value.O.(x * (x + Value.of_float 3.)) let%expect_test "jit'" = (* Suppresses noisy XLA log message *) Core_unix.putenv ~key:"TF_CPP_MIN_LOG_LEVEL" ~data:"2"; - jit' ~print_hlo:true ~f:foo (Value.of_float 2.) |> [%sexp_of: Value.t] |> print_s; + Staged.unstage (jit' ~print_hlo:true ~f:foo ()) (Value.of_float 2.) + |> [%sexp_of: Value.t] + |> print_s; [%expect {| - HloModule xla_call.6, entry_computation_layout={(f64[])->(f64[])} + HloModule xla_call.6, entry_computation_layout={(f64[],f64[])->(f64[])} - ENTRY %xla_call.6 (v_0.1: f64[]) -> (f64[]) { + ENTRY %xla_call.6 (v_0.1: f64[], c_0.2: f64[]) -> (f64[]) { %v_0.1 = f64[] parameter(0) - %constant.2 = f64[] constant(3) - %add.3 = f64[] add(f64[] %v_0.1, f64[] %constant.2) + %c_0.2 = f64[] parameter(1) + %add.3 = f64[] add(f64[] %v_0.1, f64[] %c_0.2) %multiply.4 = f64[] multiply(f64[] %v_0.1, f64[] %add.3) ROOT %tuple.5 = (f64[]) tuple(f64[] %multiply.4) } (Tensor 10 Float) |}]; (* Two-argument function *) - jit - ~print_hlo:true - (module Treeable.Tuple2 (Value) (Value)) - (module Value) - ~f:(fun (a, b) -> Value.O.(Value.sin a * Value.cos b)) + Staged.unstage + (jit + ~print_hlo:true + (module Treeable.Tuple2 (Value) (Value)) + (module Value) + ~f:(fun (a, b) -> Value.O.(Value.sin a * Value.cos b)) + ()) (Value.of_float 2., Value.of_float 3.) |> [%sexp_of: Value.t] |> print_s; @@ -225,11 +277,13 @@ let%expect_test "jit'" = let%expect_test "jit and matmul" = let a = Tensor.of_list2_exn Float [ [ 1.; 2. ]; [ 3.; 4. ] ] |> Value.of_tensor in let b = Tensor.of_list2_exn Float [ [ 5.; 6. ]; [ 7.; 8. ] ] |> Value.of_tensor in - jit - ~print_hlo:true - (module Treeable.Tuple2 (Value) (Value)) - (module Value) - ~f:(fun (a, b) -> Value.matmul a b) + Staged.unstage + (jit + ~print_hlo:true + (module Treeable.Tuple2 (Value) (Value)) + (module Value) + ~f:(fun (a, b) -> Value.matmul a b) + ()) (a, b) |> [%sexp_of: Value.t] |> print_s; @@ -248,9 +302,8 @@ let%expect_test "jit and matmul" = ;; let%expect_test "jit and sum" = - jit' - ~print_hlo:true - ~f:(fun x -> Value.sum x) + Staged.unstage + (jit' ~print_hlo:true ~f:(fun x -> Value.sum x) ()) (Value.of_tensor (Tensor.of_list2_exn Float [ [ 1.; 2. ]; [ 3.; 4. ] ])) |> [%sexp_of: Value.t] |> print_s; @@ -275,9 +328,8 @@ let%expect_test "jit and sum" = ;; let%expect_test "jit and broadcast" = - jit' - ~print_hlo:true - ~f:(fun x -> Value.broadcast x ~dims:[| 2; 2; 2 |]) + Staged.unstage + (jit' ~print_hlo:true ~f:(fun x -> Value.broadcast x ~dims:[: 2; 2; 2 :]) ()) (Value.of_tensor (Tensor.of_list2_exn Float [ [ 1.; 2. ]; [ 3.; 4. ] ])) |> [%sexp_of: Value.t] |> print_s; @@ -293,3 +345,113 @@ let%expect_test "jit and broadcast" = (Tensor (((1 2) (3 4)) ((1 2) (3 4))) (dims (2 2 2)) (type_ Float)) |}] ;; + +let%expect_test "jit caches the compiled executable per input structure" = + let jitted = Staged.unstage (jit' ~print_hlo:true ~f:(fun x -> Value.O.(x * x)) ()) in + (* The first call traces and compiles, printing the HLO. *) + let a = jitted (Value.of_float 2.) in + [%expect + {| + HloModule xla_call.35, entry_computation_layout={(f64[])->(f64[])} + + ENTRY %xla_call.35 (v_0.32: f64[]) -> (f64[]) { + %v_0.32 = f64[] parameter(0) + %multiply.33 = f64[] multiply(f64[] %v_0.32, f64[] %v_0.32) + ROOT %tuple.34 = (f64[]) tuple(f64[] %multiply.33) + } + |}]; + (* A second call with the same input structure reuses the compiled executable - nothing + is printed. *) + let b = jitted (Value.of_float 5.) in + [%expect {| |}]; + (* A different input structure is a separate cache entry, so it compiles again. *) + let c = jitted (Value.of_tensor (Tensor.of_list Float [ 1.; 2. ])) in + [%expect + {| + HloModule xla_call.39, entry_computation_layout={(f64[2]{0})->(f64[2]{0})} + + ENTRY %xla_call.39 (v_0.36: f64[2]) -> (f64[2]) { + %v_0.36 = f64[2]{0} parameter(0) + %multiply.37 = f64[2]{0} multiply(f64[2]{0} %v_0.36, f64[2]{0} %v_0.36) + ROOT %tuple.38 = (f64[2]{0}) tuple(f64[2]{0} %multiply.37) + } + |}]; + print_s [%message (a : Value.t) (b : Value.t) (c : Value.t)]; + [%expect + {| + ((a (Tensor 4 Float)) (b (Tensor 25 Float)) + (c (Tensor (1 4) (dims (2)) (type_ Float)))) + |}] +;; + +let%expect_test "distinct constants are fed as separate runtime parameters" = + (* [x * 3 + 5] hoists two constants, so the HLO has two const parameters. Each must + receive its own value: for x = 2 the result is 11, whereas feeding the constants in a + swapped order ([x * 5 + 3]) would give 13. *) + Staged.unstage + (jit' + ~print_hlo:true + ~f:(fun x -> Value.O.((x * Value.of_float 3.) + Value.of_float 5.)) + ()) + (Value.of_float 2.) + |> [%sexp_of: Value.t] + |> print_s; + [%expect + {| + HloModule xla_call.46, entry_computation_layout={(f64[],f64[],f64[])->(f64[])} + + ENTRY %xla_call.46 (v_0.40: f64[], c_0.41: f64[], c_1.42: f64[]) -> (f64[]) { + %v_0.40 = f64[] parameter(0) + %c_0.41 = f64[] parameter(1) + %multiply.43 = f64[] multiply(f64[] %v_0.40, f64[] %c_0.41) + %c_1.42 = f64[] parameter(2) + %add.44 = f64[] add(f64[] %multiply.43, f64[] %c_1.42) + ROOT %tuple.45 = (f64[]) tuple(f64[] %add.44) + } + (Tensor 11 Float) + |}] +;; + +let%expect_test "one executable serves a closure whose captured constant changes" = + let c = ref 3. in + let jitted = + Staged.unstage (jit' ~print_hlo:true ~f:(fun x -> Value.O.(x * Value.of_float !c)) ()) + in + (* The first call compiles (HLO printed), feeding the current constant as a parameter. *) + let r1 = jitted (Value.of_float 2.) in + [%expect + {| + HloModule xla_call.51, entry_computation_layout={(f64[],f64[])->(f64[])} + + ENTRY %xla_call.51 (v_0.47: f64[], c_0.48: f64[]) -> (f64[]) { + %v_0.47 = f64[] parameter(0) + %c_0.48 = f64[] parameter(1) + %multiply.49 = f64[] multiply(f64[] %v_0.47, f64[] %c_0.48) + ROOT %tuple.50 = (f64[]) tuple(f64[] %multiply.49) + } + |}]; + (* The captured constant changes, but the structure does not: the executable is reused + (no HLO) and the new value is fed in - so r1 uses 3 (= 6) and r2 uses 10 (= 20). *) + c := 10.; + let r2 = jitted (Value.of_float 2.) in + [%expect {| |}]; + print_s [%message (r1 : Value.t) (r2 : Value.t)]; + [%expect {| ((r1 (Tensor 6 Float)) (r2 (Tensor 20 Float))) |}] +;; + +let%expect_test "a cache hit feeds the call's own (multiple) constants, in order" = + let c1 = ref 3. + and c2 = ref 5. in + let jitted = + Staged.unstage + (jit' ~f:(fun x -> Value.O.((x * Value.of_float !c1) + Value.of_float !c2)) ()) + in + let r1 = jitted (Value.of_float 2.) in + c1 := 10.; + c2 := 100.; + (* The reused executable is fed [c1 = 10], [c2 = 100] in the right slots: 2*10 + 100 = + 120 (a swapped order would give 2*100 + 10 = 210). *) + let r2 = jitted (Value.of_float 2.) in + print_s [%message (r1 : Value.t) (r2 : Value.t)]; + [%expect {| ((r1 (Tensor 11 Float)) (r2 (Tensor 120 Float))) |}] +;; diff --git a/test/test_eval_backends.ml b/test/test_eval_backends.ml index de0050a..da84a26 100644 --- a/test/test_eval_backends.ml +++ b/test/test_eval_backends.ml @@ -3,8 +3,10 @@ open! Fox_core open! Base_quickcheck module Dims = struct - type t = int array [@@deriving quickcheck, sexp_of, compare] + type t = int iarray [@@deriving sexp_of, compare] + (* The generator builds a valid (positive, exactly-multiplying) dims array and freezes + it; the observer and shrinker defer to the [int array] instances. *) let quickcheck_generator = Generator.fixed_point (fun quickcheck_generator -> let open Generator.Let_syntax in @@ -18,6 +20,18 @@ module Dims = struct | remaining_size -> let%map dims = Generator.with_size quickcheck_generator ~size:remaining_size in Array.append dims [| this_dim |])) + |> Generator.map ~f:Iarray.of_array + ;; + + let quickcheck_observer = + Observer.unmap [%quickcheck.observer: int array] ~f:Iarray.to_array + ;; + + let quickcheck_shrinker = + Shrinker.map + [%quickcheck.shrinker: int array] + ~f:Iarray.of_array + ~f_inverse:Iarray.to_array ;; end @@ -26,7 +40,7 @@ module Tensor = struct let quickcheck_generator_with_dims ~dims = let open Generator.Let_syntax in - let total_elements = Array.fold dims ~init:1 ~f:( * ) in + let total_elements = Iarray.fold dims ~init:1 ~f:( * ) in let%map values = List.init total_elements ~f:(fun _ -> Generator.float_inclusive (-100.) 100.) |> Generator.all @@ -44,10 +58,10 @@ module Tensor = struct Shrinker.create (fun (Tensor.T tensor) -> let dims = Tensor.Typed.dims tensor in let sliced_tensors = - if Array.length dims > 0 + if Iarray.length dims > 0 then - Sequence.init dims.(0) ~f:(fun i -> - T (Tensor.Typed.left_slice tensor ~indices:[| i |])) + Sequence.init dims.:(0) ~f:(fun i -> + T (Tensor.Typed.left_slice tensor ~indices:[: i :])) else Sequence.empty in let smaller_tensors = @@ -179,7 +193,7 @@ let op_generator ~values_by_shape = let%bind.Option all_shapes = match Map.keys values_by_shape - |> List.filter ~f:(fun shape -> Array.length (Shape.dims shape) = 2) + |> List.filter ~f:(fun shape -> Iarray.length (Shape.dims shape) = 2) with | [] -> None | all_shapes -> Some all_shapes @@ -188,8 +202,8 @@ let op_generator ~values_by_shape = List.concat_map all_shapes ~f:(fun lhs_shape -> List.filter_map all_shapes ~f:(fun rhs_shape -> match lhs_shape, rhs_shape with - | { dims = [| _; m |]; type_ = _ }, { dims = [| m' |]; type_ = _ } - | { dims = [| _; m |]; type_ = _ }, { dims = [| m'; _ |]; type_ = _ } -> + | { dims = [: _; m :]; type_ = _ }, { dims = [: m' :]; type_ = _ } + | { dims = [: _; m :]; type_ = _ }, { dims = [: m'; _ :]; type_ = _ } -> Option.some_if (m = m') (lhs_shape, rhs_shape) | _ -> None)) with @@ -205,7 +219,7 @@ let op_generator ~values_by_shape = | Transpose () -> (match Map.keys values_by_shape - |> List.filter ~f:(fun shape -> Array.length (Shape.dims shape) = 2) + |> List.filter ~f:(fun shape -> Iarray.length (Shape.dims shape) = 2) with | [] -> return None | all_dims -> @@ -225,7 +239,7 @@ let op_generator ~values_by_shape = in (match Map.keys values_by_shape - |> List.filter ~f:(fun shape -> Array.length (Shape.dims shape) > max_index) + |> List.filter ~f:(fun shape -> Iarray.length (Shape.dims shape) > max_index) with | [] -> return None | all_shapes -> @@ -270,6 +284,7 @@ let expr_generator ~op_nums = let out = List.hd_exn equations |> Expr.Eq.var in Expr.create ~parameters:[ arg ] + ~consts:Expr.Var.Map.empty ~equations:(List.rev equations) ~return_vals:[ Expr.Atom.Var out ] ~out_tree_def:(Value_tree.Def.leaf ~dims:(Expr.Var.dims out)) @@ -279,7 +294,7 @@ let%expect_test "expr_generator" = let random = Splittable_random.of_int 0 in for i = 1 to 5 do let expr = Generator.generate (expr_generator ~op_nums:i) ~size:6 ~random in - Expr.to_string_hum expr |> print_endline; + Expr.to_string_hum expr ~value_to_string:Value.to_string |> print_endline; print_endline "--------------------------------" done; [%expect @@ -333,12 +348,15 @@ let%expect_test "eval expr vs xla" = (fun_generator ~op_nums:1) ~trials:300 ~sexp_of:(fun (tensor, expr) -> - [%sexp { tensor : Tensor.t; expr : string = Expr.to_string_hum expr }]) + [%sexp + { tensor : Tensor.t + ; expr : string = Expr.to_string_hum expr ~value_to_string:Value.to_string + }]) ~f:(fun (tensor, expr) -> let f value = eval_expr' expr value in let value = Value.of_tensor tensor in - let eval_result = Eval.handle ~f:(fun () -> f value) in - let xla_result = Fox_jit.jit' ~f value in + let eval_result = eval ~f:(fun () -> f value) in + let xla_result = Staged.unstage (Fox_jit.jit' ~f ()) value in assert ( Tensor.allclose ~equal_nan:true @@ -348,14 +366,13 @@ let%expect_test "eval expr vs xla" = let%expect_test "grad+jit vs grad+eval" = let test ~f ~x = - Eval.handle ~f:(fun () -> f x) |> [%sexp_of: Value.t] |> print_s; - Fox_jit.jit' ~f x |> [%sexp_of: Value.t] |> print_s + eval ~f:(fun () -> f x) |> [%sexp_of: Value.t] |> print_s; + Staged.unstage (Fox_jit.jit' ~f ()) x |> [%sexp_of: Value.t] |> print_s in test ~f:(fun value -> grad' ~f:(fun value -> Value.O.(value * value)) ~x:value) ~x:(Value.of_float 1.); - [%expect - {| + [%expect {| (Tensor 2 Float) (Tensor 2 Float) |}]; @@ -382,12 +399,15 @@ let%expect_test "eval grad expr vs xla" = Core_unix.putenv ~key:"TF_CPP_MIN_LOG_LEVEL" ~data:"2"; Quickcheck.test (fun_generator ~op_nums:1) - (* TODO: without a periodic [Gc.full_major ()], pthread_create fails with EAGAIN - and causes a SIGABRT (at least on macos) when ~trials is set to more than 300. - This is likely a result of not properly releasing resources somewhere. *) + (* TODO: without a periodic [Gc.full_major ()], pthread_create fails with EAGAIN and + causes a SIGABRT (at least on macos) when ~trials is set to more than 300. This is + likely a result of not properly releasing resources somewhere. *) ~trials:200 ~sexp_of:(fun (tensor, expr) -> - [%sexp { tensor : Tensor.t; expr : string = Expr.to_string_hum expr }]) + [%sexp + { tensor : Tensor.t + ; expr : string = Expr.to_string_hum expr ~value_to_string:Value.to_string + }]) ~f:(fun (tensor, expr) -> let f value = grad' @@ -398,14 +418,14 @@ let%expect_test "eval grad expr vs xla" = | T Float -> Value.sum expr_result | T Bool -> (* TODO: somehow prevent return type from being a boolean? *) - (* We can't really differentiate bool tensors, so we just sum the - input instead. *) + (* We can't really differentiate bool tensors, so we just sum the input + instead. *) Value.sum value) ~x:value in let value = Value.of_tensor tensor in - let eval_result = Eval.handle ~f:(fun () -> f value) in - let xla_result = Fox_jit.jit' ~f value in + let eval_result = eval ~f:(fun () -> f value) in + let xla_result = Staged.unstage (Fox_jit.jit' ~f ()) value in assert ( Tensor.allclose ~equal_nan:true diff --git a/test/test_fox.ml b/test/test_fox.ml index e6cfde5..b66a1ac 100644 --- a/test/test_fox.ml +++ b/test/test_fox.ml @@ -8,12 +8,12 @@ let f x = ;; let%expect_test "eval" = - Eval.handle ~f:(fun () -> f (Value.of_float 3.)) |> [%sexp_of: Value.t] |> print_s; + eval ~f:(fun () -> f (Value.of_float 3.)) |> [%sexp_of: Value.t] |> print_s; [%expect {| (Tensor 2.7177599838802657 Float) |}] ;; let%expect_test "jvp" = - Eval.handle ~f:(fun () -> + eval ~f:(fun () -> jvp' ~f:Value.sin ~primal:(Value.of_float 3.) ~tangent:(Value.of_float 1.) |> snd |> [%sexp_of: Value.t] @@ -24,12 +24,11 @@ let%expect_test "jvp" = (Tensor -0.98999249660044542 Float) (Tensor -0.98999249660044542 Float) |}]; - Eval.handle ~f:(fun () -> - jvp' ~f ~primal:(Value.of_float 3.) ~tangent:(Value.of_float 1.)) + eval ~f:(fun () -> jvp' ~f ~primal:(Value.of_float 3.) ~tangent:(Value.of_float 1.)) |> [%sexp_of: Value.t * Value.t] |> print_s; [%expect {| ((Tensor 2.7177599838802657 Float) (Tensor 2.9799849932008908 Float)) |}]; - Eval.handle ~f:(fun () -> + eval ~f:(fun () -> let deriv ~n = nth_order_derivative ~n ~f:Value.sin ~x:(Value.of_float 3.) |> [%sexp_of: Value.t] diff --git a/test/test_handler.ml b/test/test_handler.ml new file mode 100644 index 0000000..580e75d --- /dev/null +++ b/test/test_handler.ml @@ -0,0 +1,156 @@ +open! Core +open! Fox_core + +let foo x = Value.O.(x * (x + Value.of_float 3.)) + +let%expect_test "eval_expr" = + eval ~f:(fun () -> + let expr = build_expr' ~f:foo ~in_dims:[::] in + eval_expr' expr (Value.of_float 2.)) + |> [%sexp_of: Value.t] + |> print_s; + [%expect {| (Tensor 10 Float) |}] +;; + +let%expect_test "jvp and eval_expr" = + eval ~f:(fun () -> + jvp' + ~f:(fun x -> + let expr = build_expr' ~f:foo ~in_dims:[::] in + eval_expr' expr x) + ~primal:(Value.of_float 2.) + ~tangent:(Value.of_float 1.)) + |> [%sexp_of: Value.t * Value.t] + |> print_s; + [%expect {| ((Tensor 10 Float) (Tensor 7 Float)) |}] +;; + +let%expect_test "linearize" = + let y, f_lin = + eval ~f:(fun () -> linearize' ~f:Value.sin ~primals:(Value.of_float 3.)) + in + print_s [%message "" (y : Value.t) (Float.sin 3. : float)]; + [%expect + {| ((y (Tensor 0.14112000805986721 Float)) ("Float.sin 3." 0.14112000805986721)) |}]; + let y' = eval ~f:(fun () -> f_lin (Value.of_float 1.)) in + print_s [%message "" (y' : Value.t) (Float.cos 3. : float)]; + [%expect + {| + ((y' (Tensor -0.98999249660044542 Float)) + ("Float.cos 3." -0.98999249660044542)) + |}]; + let y, f_lin = + eval ~f:(fun () -> + linearize' + ~f:(fun x -> + let y = Value.O.(Value.sin x * Value.of_float 2.) in + Value.O.(-y + x)) + ~primals:(Value.of_float 3.)) + in + let y' = eval ~f:(fun () -> f_lin (Value.of_float 1.)) in + print_s [%message "" (y : Value.t) (y' : Value.t)]; + [%expect + {| + ((y (Tensor 2.7177599838802657 Float)) + (y' (Tensor 2.9799849932008908 Float))) + |}]; + let f a = + let b = Value.sin a in + let c = Value.neg b in + c + in + let expr = build_expr' ~f ~in_dims:[::] in + Expr.to_string_hum expr ~value_to_string:Value.to_string |> print_endline; + [%expect + {| + v_0[]: float -> + v_1[]: float = sin v_0; + v_2[]: float = neg v_1; + ( v_2 ) + |}]; + let expr = + build_expr + (module Value.Tuple2) + (module Value.Tuple2) + ~f:(fun (a, b) -> jvp' ~f ~primal:a ~tangent:b) + ~in_tree_def:(Value.Tuple2.tree_def ~dims1:[::] ~dims2:[::]) + in + Expr.to_string_hum expr ~value_to_string:Value.to_string |> print_endline; + [%expect + {| + v_0[]: float v_1[]: float -> + v_2[]: float = cos v_0; + v_3[]: float = mul v_2 v_1; + v_4[]: float = sin v_0; + v_5[]: float = neg v_3; + v_6[]: float = neg v_4; + ( v_6, v_5 ) + |}]; + let expr = + build_expr' + ~f:(fun x -> + let y, _f_lin = linearize' ~f ~primals:x in + y) + ~in_dims:[::] + in + Expr.to_string_hum expr ~value_to_string:Value.to_string |> print_endline; + [%expect + {| + v_0[]: float -> + v_1[]: float = cos v_0; + v_2[]: float = sin v_0; + v_3[]: float = neg v_2; + ( v_3 ) + |}]; + let _y, f_lin = eval ~f:(fun () -> linearize' ~f ~primals:(Value.of_float 0.)) in + let expr = build_expr' ~f:f_lin ~in_dims:[::] in + Expr.to_string_hum expr ~value_to_string:Value.to_string |> print_endline; + [%expect + {| + v_0[]: float -> + consts: + c_0[]: float = (Tensor 1 Float) + v_1[]: float = mul c_0 v_0; + v_2[]: float = neg v_1; + ( v_2 ) + |}] +;; + +let%expect_test "grad" = + let y, f_vjp = eval ~f:(fun () -> vjp' ~f:Value.sin ~primal:(Value.of_float 3.)) in + let y' = eval ~f:(fun () -> f_vjp (Value.of_float 1.)) in + print_s [%message "" (y : Value.t) (y' : Value.t)]; + [%expect + {| + ((y (Tensor 0.14112000805986721 Float)) + (y' (Tensor -0.98999249660044542 Float))) + |}]; + eval ~f:(fun () -> grad' ~f:(fun x -> Value.O.(x * x)) ~x:(Value.of_float 3.)) + |> [%sexp_of: Value.t] + |> print_s; + [%expect {| (Tensor 6 Float) |}]; + eval ~f:(fun () -> + grad' + ~f:(fun x -> + let y = Value.O.(Value.sin x * Value.of_float 2.) in + Value.O.(-y + x)) + ~x:(Value.of_float 3.)) + |> [%sexp_of: Value.t] + |> print_s; + [%expect {| (Tensor 2.9799849932008908 Float) |}]; + eval ~f:(fun () -> + grad' + ~f:(Value.sum ~keep_dims:false) + ~x:(Value.of_tensor (Tensor.of_list2_exn Float [ [ 1.; 2. ]; [ 3.; 4. ] ]))) + |> [%sexp_of: Value.t] + |> print_s; + [%expect {| (Tensor ((1 1) (1 1)) (dims (2 2)) (type_ Float)) |}]; + eval ~f:(fun () -> + grad' + ~f:(fun x -> + Value.broadcast x ~dims:[: 3; 4 :] |> Value.sum ~dims:(`Just [ 1 ]) |> Value.mean) + ~x:(Value.of_typed_tensor (Tensor.Typed.arange 4))) + |> [%sexp_of: Value.t] + |> print_s; + [%expect {| (Tensor (1 1 1 1) (dims (4)) (type_ Float)) |}] +;; diff --git a/test/test_jvp.ml b/test/test_jvp.ml new file mode 100644 index 0000000..7c6091e --- /dev/null +++ b/test/test_jvp.ml @@ -0,0 +1,74 @@ +open! Core +open! Fox_core + +(* - foo(x) = x(x + 3) = x^2 + 3x + - foo'(x) = 2x + 3 + - foo''(x) = 2 + - foo'''(x) = 0 +*) +let foo x = Value.O.(x * (x + Value.of_float 3.)) + +let%expect_test "foo" = + eval ~f:(fun () -> foo (Value.of_float 2.)) |> [%sexp_of: Value.t] |> print_s; + [%expect {| (Tensor 10 Float) |}] +;; + +let%expect_test "jvp'" = + eval ~f:(fun () -> + jvp + (module Value) + (module Value) + ~f:foo + ~primals:(Value.of_float 2.) + ~tangents:(Value.of_float 1.)) + |> [%sexp_of: Value.t * Value.t] + |> print_s; + [%expect {| ((Tensor 10 Float) (Tensor 7 Float)) |}] +;; + +let%expect_test "jvp" = + eval ~f:(fun () -> jvp' ~f:foo ~primal:(Value.of_float 2.) ~tangent:(Value.of_float 1.)) + |> [%sexp_of: Value.t * Value.t] + |> print_s; + [%expect {| ((Tensor 10 Float) (Tensor 7 Float)) |}]; + eval ~f:(fun () -> + jvp' + ~f:(fun x -> + let _, tangent = jvp' ~f:foo ~primal:x ~tangent:(Value.of_float 1.) in + tangent) + ~primal:(Value.of_float 2.) + ~tangent:(Value.of_float 1.)) + |> [%sexp_of: Value.t * Value.t] + |> print_s; + [%expect {| ((Tensor 7 Float) (Tensor 2 Float)) |}] +;; + +let%expect_test "nth_order_derivative" = + let print ~n = + eval ~f:(fun () -> nth_order_derivative ~n ~f:foo ~x:(Value.of_float 2.)) + |> [%sexp_of: Value.t] + |> print_s + in + print ~n:0; + [%expect {| (Tensor 10 Float) |}]; + print ~n:1; + [%expect {| (Tensor 7 Float) |}]; + print ~n:2; + [%expect {| (Tensor 2 Float) |}]; + print ~n:3; + [%expect {| (Tensor 0 Float) |}]; + print ~n:4; + [%expect {| (Tensor 0 Float) |}] +;; + +let%expect_test "pertubation confusion avoidance" = + let f x = + let g (_y : Value.t) = x in + let should_be_zero = derivative ~f:g ~x:(Value.of_float 0.) in + Value.O.(x * should_be_zero) + in + eval ~f:(fun () -> derivative ~f ~x:(Value.of_float 0.)) + |> [%sexp_of: Value.t] + |> print_s; + [%expect {| (Tensor 0 Float) |}] +;; diff --git a/test/test_linear_regression.ml b/test/test_linear_regression.ml index e282a6f..46dae92 100644 --- a/test/test_linear_regression.ml +++ b/test/test_linear_regression.ml @@ -7,59 +7,67 @@ let%expect_test "linear regression" = let num_iters = 100 in let learning_rate = 0.01 in let rng = Splittable_random.of_int 0 in - let x = Tensor.normal ~dims:[| num_examples; num_features |] ~rng () in - let true_params = Tensor.normal ~dims:[| num_features |] ~rng () in + let x = Tensor.normal ~dims:[: num_examples; num_features :] ~rng () in + let true_params = Tensor.normal ~dims:[: num_features :] ~rng () in let y = Tensor.O.( Tensor.matmul x true_params - + Tensor.normal ~dims:[| num_examples |] ~std:1e-2 ~rng ()) + + Tensor.normal ~dims:[: num_examples :] ~std:1e-2 ~rng ()) in let loss weights = let open Value.O in let error = Value.of_tensor y - Value.matmul (Value.of_tensor x) weights in Value.mean (error * error) in - let weights = ref (Tensor.normal ~dims:[| num_features |] ~rng ()) in + let weights = ref (Tensor.normal ~dims:[: num_features :] ~rng ()) in let loss_grad x = grad' ~f:loss ~x in - build_expr' ~f:loss ~in_dims:[| num_features |] |> Expr.to_string_hum |> print_endline; + let expr = build_expr' ~f:loss ~in_dims:[: num_features :] in + Expr.to_string_hum expr ~value_to_string:Value.to_string |> print_endline; [%expect {| v_0[2]: float -> - v_1[10]: float = matmul (Tensor((0.39995642633665462 1.1602368073797789)(1.1461698444484156 -0.27508260258159217)(-0.79429206930935448 -0.79691620100672833)(0.37183586790086887 0.08252530331332697)(-0.640823054843121 -0.972395618004264)(-0.24531465496199806 0.32943595058180369)(-0.49745709501288782 -0.5186552221834726)(-0.49202679675483579 -0.30011198307839476)(-2.2219543337312779 0.66654181134893775)(0.458720357930648 0.0010728286252209105))(dims(10 2))(type_ Float)) v_0; - v_2[10]: float = sub (Tensor(-0.0064666378598043162 1.9665677929469698 -0.82202018980048819 0.55321043273971726 -0.47887754772981445 -0.59419529415798567 -0.479101132934855 -0.60615974619803537 -3.888855769581165 0.72374784193750175)(dims(10))(type_ Float)) v_1; + consts: + c_0[10,2]: float = (Tensor((0.39995642633665462 1.1602368073797789)(1.1461698444484156 -0.27508260258159217)(-0.79429206930935425 -0.79691620100672811)(0.37183586790086887 0.08252530331332697)(-0.640823054843121 -0.972395618004264)(-0.24531465496199806 0.32943595058180369)(-0.49745709501288782 -0.5186552221834726)(-0.49202679675483579 -0.30011198307839476)(-2.2219543337312779 0.66654181134893775)(0.458720357930648 0.0010728286252209105))(dims(10 2))(type_ Float)) + c_1[10]: float = (Tensor(-0.0064666378598043162 1.9665677929469698 -0.82202018980048785 0.55321043273971726 -0.47887754772981445 -0.59419529415798567 -0.479101132934855 -0.60615974619803537 -3.888855769581165 0.72374784193750175)(dims(10))(type_ Float)) + c_2[]: float = (Tensor 10 Float) + v_1[10]: float = matmul c_0 v_0; + v_2[10]: float = sub c_1 v_1; v_3[10]: float = mul v_2 v_2; v_4[]: float = sum v_3 dims=all keep_dims=false; - v_5[]: float = broadcast (Tensor 10 Float) dims=[]; + v_5[]: float = broadcast c_2 dims=[]; v_6[]: float = div v_4 v_5; ( v_6 ) |}]; - let _y, f_jvp = - Eval.handle ~f:(fun () -> vjp' ~f:loss ~primal:(Value.of_tensor !weights)) - in - build_expr' ~f:f_jvp ~in_dims:[||] |> Expr.to_string_hum |> print_endline; + let _y, f_jvp = eval ~f:(fun () -> vjp' ~f:loss ~primal:(Value.of_tensor !weights)) in + let expr = build_expr' ~f:f_jvp ~in_dims:[::] in + Expr.to_string_hum expr ~value_to_string:Value.to_string |> print_endline; [%expect {| v_0[]: float -> - v_1[]: float = div v_0 (Tensor 100 Float); - v_2[]: float = mul (Tensor 10 Float) v_1; + consts: + c_0[]: float = (Tensor 100 Float) + c_1[]: float = (Tensor 10 Float) + c_2[10]: float = (Tensor(-0.22987507367862317 4.3463761730905182 -1.6427703619487455 1.1811642145933288 -0.86466781807436588 -1.3305171979260781 -0.976724348558953 -1.2769362756602236 -8.614129411878265 1.5829382308162989)(dims(10))(type_ Float)) + c_3[10,2]: float = (Tensor((0.39995642633665462 1.1602368073797789)(1.1461698444484156 -0.27508260258159217)(-0.79429206930935425 -0.79691620100672811)(0.37183586790086887 0.08252530331332697)(-0.640823054843121 -0.972395618004264)(-0.24531465496199806 0.32943595058180369)(-0.49745709501288782 -0.5186552221834726)(-0.49202679675483579 -0.30011198307839476)(-2.2219543337312779 0.66654181134893775)(0.458720357930648 0.0010728286252209105))(dims(10 2))(type_ Float)) + v_1[]: float = div v_0 c_0; + v_2[]: float = mul c_1 v_1; v_3[1]: float = reshape v_2 dims=[1]; v_4[10]: float = broadcast v_3 dims=[10]; - v_5[10]: float = mul (Tensor(-0.22987507367862317 4.3463761730905182 -1.6427703619487461 1.1811642145933288 -0.86466781807436588 -1.3305171979260781 -0.976724348558953 -1.2769362756602236 -8.614129411878265 1.5829382308162989)(dims(10))(type_ Float)) v_4; - v_6[10]: float = mul (Tensor(-0.22987507367862317 4.3463761730905182 -1.6427703619487461 1.1811642145933288 -0.86466781807436588 -1.3305171979260781 -0.976724348558953 -1.2769362756602236 -8.614129411878265 1.5829382308162989)(dims(10))(type_ Float)) v_4; + v_5[10]: float = mul c_2 v_4; + v_6[10]: float = mul c_2 v_4; v_7[10]: float = add v_5 v_6; v_8[10]: float = neg v_7; - v_9[2,10]: float = transpose (Tensor((0.39995642633665462 1.1602368073797789)(1.1461698444484156 -0.27508260258159217)(-0.79429206930935448 -0.79691620100672833)(0.37183586790086887 0.08252530331332697)(-0.640823054843121 -0.972395618004264)(-0.24531465496199806 0.32943595058180369)(-0.49745709501288782 -0.5186552221834726)(-0.49202679675483579 -0.30011198307839476)(-2.2219543337312779 0.66654181134893775)(0.458720357930648 0.0010728286252209105))(dims(10 2))(type_ Float)); + v_9[2,10]: float = transpose c_3; v_10[2]: float = matmul v_9 v_8; ( v_10 ) |}]; for i = 0 to num_iters do let grads = - Eval.handle ~f:(fun () -> loss_grad (Value.of_tensor !weights)) - |> Value.to_tensor_exn + eval ~f:(fun () -> loss_grad (Value.of_tensor !weights)) |> Value.to_tensor_exn in weights := Tensor.sub !weights (Tensor.scale grads learning_rate); let loss = - Eval.handle ~f:(fun () -> loss (Value.of_tensor !weights)) |> Value.to_float_exn + eval ~f:(fun () -> loss (Value.of_tensor !weights)) |> Value.to_float_exn in if i mod 10 = 0 then print_s [%message "" (i : int) (loss : float)] done; @@ -73,8 +81,8 @@ let%expect_test "linear regression" = ((i 50) (loss 2.1485718499399744)) ((i 60) (loss 1.5973678981829491)) ((i 70) (loss 1.1947496743966037)) - ((i 80) (loss 0.89950194008054518)) - ((i 90) (loss 0.68202100327261561)) + ((i 80) (loss 0.8995019400805454)) + ((i 90) (loss 0.68202100327261594)) ((i 100) (loss 0.52101656302408539)) |}] ;; diff --git a/test/test_partial_eval.ml b/test/test_partial_eval.ml new file mode 100644 index 0000000..2857bd0 --- /dev/null +++ b/test/test_partial_eval.ml @@ -0,0 +1,37 @@ +open! Core +open! Fox_core +open For_testing + +let%expect_test "partially_apply_expr_flat" = + let partial_values, expr = + eval ~f:(fun () -> + partially_apply_expr_flat + [ Known (Value.of_float 2.) + ; Unknown { name = "x"; shape = { dims = [::]; type_ = T Float } } + ] + ~f:(function + | [ x; y ] -> + let x2 = Value.O.(x * x) in + ( [ x2; Value.O.((x2 * y) + Value.of_float 3.); x; Value.O.((y * y) + x2) ] + , Value.tree_def ~dims:[::] ) + | _ -> assert false)) + in + print_s ([%sexp_of: Partial_value.t list] partial_values); + [%expect + {| + ((Known (Tensor 4 Float)) + (Unknown ((name p_3) (shape ((dims ()) (type_ Float))))) + (Known (Tensor 2 Float)) + (Unknown ((name p_1) (shape ((dims ()) (type_ Float)))))) + |}]; + Expr.to_string_hum expr ~value_to_string:Value.to_string |> print_endline; + [%expect + {| + x[]: float -> + p_0[]: float = mul x x; + p_1[]: float = add p_0 (Tensor 4 Float); + p_2[]: float = mul (Tensor 4 Float) x; + p_3[]: float = add p_2 (Tensor 3 Float); + ( p_3, p_1 ) + |}] +;; diff --git a/test/test_staging.ml b/test/test_staging.ml new file mode 100644 index 0000000..a3781b3 --- /dev/null +++ b/test/test_staging.ml @@ -0,0 +1,166 @@ +open! Core +open! Fox_core + +let foo x = Value.O.(x * (x + Value.of_float 3.)) + +let%expect_test "build_expr" = + let expr = build_expr' ~f:foo ~in_dims:[::] in + Expr.to_string_hum expr ~value_to_string:Value.to_string |> print_endline; + [%expect + {| + v_0[]: float -> + consts: + c_0[]: float = (Tensor 3 Float) + v_1[]: float = add v_0 c_0; + v_2[]: float = mul v_0 v_1; + ( v_2 ) + |}] +;; + +let%expect_test "build_expr2" = + (* Distinct [of_float 2.] calls are distinct values, so they are not shared. *) + let expr = + build_expr' + ~f:(fun _x -> Value.O.(Value.of_float 2. * Value.of_float 2.)) + ~in_dims:[::] + in + Expr.to_string_hum expr ~value_to_string:Value.to_string |> print_endline; + [%expect + {| + v_0[]: float -> + consts: + c_0[]: float = (Tensor 2 Float) + c_1[]: float = (Tensor 2 Float) + v_1[]: float = mul c_1 c_0; + ( v_1 ) + |}] +;; + +let%expect_test "[Expr.map ~f:(fun _ -> ())] is a value-free structural key" = + let structure f = build_expr' ~f ~in_dims:[::] |> Expr.map ~f:(fun _ -> ()) in + let a = structure (fun x -> Value.O.((x * Value.of_float 3.) + Value.of_float 5.)) in + let b = structure (fun x -> Value.O.((x * Value.of_float 7.) + Value.of_float 9.)) in + let c = structure (fun x -> Value.O.(x * x)) in + (* The constants are erased to (), so [a] (constants 3, 5) and [b] (constants 7, 9) are + structurally equal, while [c] is a different program. *) + print_endline (Expr.to_string_hum a ~value_to_string:Unit.to_string); + [%expect + {| + v_0[]: float -> + consts: + c_0[]: float = () + c_1[]: float = () + v_1[]: float = mul v_0 c_0; + v_2[]: float = add v_1 c_1; + ( v_2 ) + |}]; + let equal = [%compare.equal: unit Expr.t] in + let same_hash x y = [%hash: unit Expr.t] x = [%hash: unit Expr.t] y in + print_s [%message (equal a b : bool) (same_hash a b : bool) (equal a c : bool)]; + [%expect {| (("equal a b" true) ("same_hash a b" true) ("equal a c" false)) |}] +;; + +let%expect_test "a constant returned directly is hoisted into consts" = + (* The returned constant occurs only in the return values, not in any equation, so + [build_expr] must resolve the return values (which is what hoists the constant) + before snapshotting [consts] - otherwise [c_0] would be referenced without appearing + in the consts map. *) + let expr = build_expr' ~f:(fun _x -> Value.of_float 5.) ~in_dims:[::] in + Expr.to_string_hum expr ~value_to_string:Value.to_string |> print_endline; + [%expect + {| + v_0[]: float -> + consts: + c_0[]: float = (Tensor 5 Float) + + ( c_0 ) + |}] +;; + +let%expect_test "shared constants are hoisted and deduplicated" = + (* [three] is a single value reused in two ops; it is hoisted into one shared const var + [c_0] rather than being embedded inline at each use. *) + let three = Value.of_float 3. in + let expr = + build_expr' ~f:(fun x -> Value.O.((x + three) * (x + three))) ~in_dims:[::] + in + Expr.to_string_hum expr ~value_to_string:Value.to_string |> print_endline; + [%expect + {| + v_0[]: float -> + consts: + c_0[]: float = (Tensor 3 Float) + v_1[]: float = add v_0 c_0; + v_2[]: float = add v_0 c_0; + v_3[]: float = mul v_2 v_1; + ( v_3 ) + |}] +;; + +let%expect_test "nth_order_derivative build_expr" = + let print ~n = + let expr = + build_expr' ~f:(fun x -> nth_order_derivative ~n ~f:foo ~x) ~in_dims:[::] + in + Expr.to_string_hum expr ~value_to_string:Value.to_string |> print_endline + in + print ~n:0; + [%expect + {| + v_0[]: float -> + consts: + c_0[]: float = (Tensor 3 Float) + v_1[]: float = add v_0 c_0; + v_2[]: float = mul v_0 v_1; + ( v_2 ) + |}]; + print ~n:1; + [%expect + {| + v_0[]: float -> + consts: + c_0[]: float = (Tensor 0 Float) + c_1[]: float = (Tensor 1 Float) + c_2[]: float = (Tensor 3 Float) + v_1[]: float = add c_1 c_0; + v_2[]: float = add v_0 c_2; + v_3[]: float = mul v_0 v_1; + v_4[]: float = mul c_1 v_2; + v_5[]: float = add v_4 v_3; + v_6[]: float = mul v_0 v_2; + ( v_5 ) + |}]; + print ~n:2; + [%expect + {| + v_0[]: float -> + consts: + c_0[]: float = (Tensor 0 Float) + c_1[]: float = (Tensor 0 Float) + c_2[]: float = (Tensor 0 Float) + c_3[]: float = (Tensor 1 Float) + c_4[]: float = (Tensor 0 Float) + c_5[]: float = (Tensor 1 Float) + c_6[]: float = (Tensor 3 Float) + c_7[]: float = (Tensor 0 Float) + v_1[]: float = add c_1 c_0; + v_2[]: float = add c_3 c_2; + v_3[]: float = add c_5 c_4; + v_4[]: float = add v_0 c_6; + v_5[]: float = mul v_0 v_1; + v_6[]: float = mul c_5 v_2; + v_7[]: float = add v_6 v_5; + v_8[]: float = mul v_0 v_2; + v_9[]: float = mul c_3 v_3; + v_10[]: float = mul c_7 v_4; + v_11[]: float = add v_10 v_9; + v_12[]: float = mul c_3 v_4; + v_13[]: float = add v_11 v_7; + v_14[]: float = add v_12 v_8; + v_15[]: float = mul v_0 v_3; + v_16[]: float = mul c_5 v_4; + v_17[]: float = add v_16 v_15; + v_18[]: float = mul v_0 v_4; + ( v_13 ) + |}] +;;