jax

c12541d5 - Refactor the jax_to_tf tests to separate the primitive test harness (#3376)

Refactor the jax_to_tf tests to separate the primitive test harness (#3376) * Refactor the jax_to_tf tests to separate the primitive test harness from the test. The goal is to have a collection of test harnesses for the JAX primitives to be able to test various implementation (JAX, NumPy, TensorFlow). For now we use these harnesses only in the jax_to_tf tests, although we can later use them for lax_test. Demonstrate the use of the harness for lax.pad and lax.squeeze, both in tf_ops_test and lax_test. The plan is to add support for more primitives as we make progress testing jax_to_tf. * Expanded pad harness with negative pads