jax

f7a44523 - Add some type helpers to lax_numpy. (#1593)

Add some type helpers to lax_numpy. (#1593) Prefer to use jax.numpy type helpers rather than numpy type helpers in various places. Cleanup in preparation for adding bfloat16 support to jax.