Accelerated machine learning with dependent types
spidr is a research project, expect breaking changes. We have yet to implement automatic differentiation and gradient-based optimizers.
See the online reference for API documentation, and the tutorials for extended discussion of spidr's architecture. In particular, make sure to read Nuisances in the Tensor API.
Installation requires curl and pack. To install spidr, install a PJRT plugin. A plugin executes a spidr program, and determines what hardware your program will run on. You can install the CPU plugin, on Linux or Mac with Apple silicon, with
pack install pjrt-plugin-xla-cpu
or read the plugin documentation for the CUDA-enabled GPU plugin and custom plugin builds.
We made spidr to try out modern programming language capabilities in machine learning systems. To this end, we chose Idris for the API; Idris is a general-purpose purely functional programming language with a particularly expressive type system. We also wanted to build something performant enough for working machine learning practitioners. Implementing efficient low-level linear algebra is not one of the project goals, so we opted to use OpenXLA's PJRT, an abstract interface for machine learning compilers for hardware accelerators.
If your spidr program compiles, and your hardware can run it, then it will run. This is primarily because Idris checks tensor shapes during compilation. For example, this will compile
x : Tensor [3] S32
x = tensor [1, 2, 3]
y : Tensor [3] S32
y = x + tensor [0, 1, 2]but this won't
failing "elaboration"
z : Tensor [3] S32
z = x + tensor [0, 1]because you can't add a vector of length two to a vector of length three. Shape manipulation extends beyond comparing literal dimension sizes to arbitrary symbolic manipulation
append : Tensor [m, p] F64 -> Tensor [n, p] F64 -> Tensor [m + n, p] F64
append x y = concat 0 x yAs a bonus, spidr programs are reproducible. Any one graph will always produce the same result when run on the same hardware.
You can run spidr programs on any hardware for which there's a PJRT plugin. CPU and CUDA plugins are provided out of the box. You can also create your own plugins with minimal code, see the guide for instructions. The libraries required to build a plugin exist for ROCm-enabled GPUs, specialised machine learning accelerators, and more.
Each PJRT plugin contains a graph compiler, and there are several compilers available. The plugins we provide out of the box use XLA, which offers substantial performance benefits via e.g. CSE and operator fusion.
This is a high-priority feature but is not yet implemented. spidr can generate new tensor graphs from existing ones. We plan to use this to implement vectorization and automatic differentiation, like JAX's vmap and grad.
I'd like to thank the Idris community for their frequent guidance and Idris itself, the Numerical Elixir team for their early binaries, Secondmind colleagues for discussions around machine learning design, friends and family for their support, Google and OpenXLA for the compiler stack, and Github for hosting.
To ask for new features or to report bugs, make a new GitHub issue. For any other questions or comments, message @joelb on the Idris community discord.