Cleanup: fix type issues in lax_numpy.py (#3816)
Cleanup: fix type issues in lax_numpy.py (#3816)

These changes are basically a no-op wth the current default types, but fixes issues if/when those types are changed to 32-bit in the future.
google/jax,tensorflow/probability,google/jax,tensorflow/probability,google/jax,google/jax