Add TPU tests for JAX and Tensorflow.#160

hertschuh merged 1 commit intokeras-team:mainkeras-team/keras-rs:mainfrom

hertschuh:tpu_testshertschuh/keras-rs:tpu_testsCopy head branch name to clipboard