jax

a1e67be7 - [jax2tf] Fix the behavior for shift when the shift amount is negative or large.

[jax2tf] Fix the behavior for shift when the shift amount is negative or large. It turns out that TF shift behavior is "implementation defined" if the shift amount is negative or larger or equal to the bitsize of the operand. This is different than XLA. We add conditionals to check and fix the corner cases, we expand the tests. A better solution would be to use an XlaShift operation. This also revealed some bugs with shift_right on TPU (for JAX). The jax2tf behavior is what we believe is the correct behavior, which means that it differs from JAX results on TPU.