jax

a07c5e11 - Fix integer overflow in np.split on windows.

Fix integer overflow in np.split on windows. Cast earlier to np.int64. This is caught by the existing random_test.py, but only is a problem on Windows which has slightly different integer type behavior.