Hyjax

pip install hyjax

Hyjax is Hy bindings for JAX.

When walking about the countryside of Italy, the people will not hesitate to tell you that JAX has “una anima di pura programmazione funzionale”. –– 🔪 JAX - The Sharp Bits 🔪

“As one would expect from its goals, artificial intelligence research generates many significant programming problems. In other programming cultures this spate of problems spawns new languages. ... We toast the Lisp programmer who pens his thoughts within nests of parentheses.” -- Alan J. Perlis

You now know enough to be dangerous with Hy. You may now smile villainously and sneak off to your Hydeaway to do unspeakable things. –– Hy Tutorial

The goals are

1. the name pun was funny
2. let fully JIT-compiled JAX code be written in idiomatic Lisp

Some examples loosely inspired by the JAX Quickstart tutorials are in src/experiments.hy. The Hy and JAX docs might also be helpful.

Features

• (mapv f vec) as vmap(f)(vec)
• (if/ja pred then else) as lax.cond
• (defn/j f [args] body) as @jit def f(args): ...
  • identical syntax to Hy's defn; supports other decorators, annotations, variadic & keyword args
• (if pred then else) inside defn/j as lax.cond
  • works with any macros that compile to if, such as cond
• binding for lax.while_loop
• binding for lax.fori_loop