xla

8a9fdcfb - tweak `atol` and `rtol` for test_simple_model_with_different_input_shape