jax

e1ea936f - Added a custom lowering rule for pow which special-cases weak dtypes