jax

e36c72b9 - Make ad_util.zero a class that carries avals (similar to UndefinedPrimal) (#3222)