jax

f036f5dd - Avoid direct type/dtype comparisons to fix NumPy 1.19 deprecation war… (#3543)

Avoid direct type/dtype comparisons to fix NumPy 1.19 deprecation war… (#3543) * Avoid direct type/dtype comparisons to fix NumPy 1.19 deprecation warnings. * Pin a newer tf-nightly to fix jax2tf tests for NumPy 1.19.0