jax

815f0615 - fix tests for random dtype