dax | A fast AI design library based on JAX.