Disable XLA with TensorFlow determinisim

[Frightera]

Currently if someone runs https://keras.io/examples/keras_recipes/reproducibility_recipes/, gets an error:

(PR for fixing recipe keras-team/keras-io#1941)

UnimplementedError: Graph execution error:

Detected at node gradient_tape/sequential_1/max_pooling2d_1_2/MaxPool2d/MaxPoolGrad defined at (most recent call last):
<stack traces unavailable>
GPU MaxPool gradient ops do not yet have a deterministic XLA implementation.
	 [[{{node gradient_tape/sequential_1/max_pooling2d_1_2/MaxPool2d/MaxPoolGrad}}]]
	tf2xla conversion failed while converting __inference_one_step_on_data_2978[]. Run with TF_DUMP_GRAPH_PREFIX=/path/to/dump/dir and --vmodule=xla_compiler=2 to obtain a dump of the compiled functions.
	 [[StatefulPartitionedCall]] [Op:__inference_one_step_on_iterator_3061]

Considering Keras 2, which had jit_compile = None (not enabled by default) it's better not to use XLA when TF determinism is enabled?

(I can check the failed tests if the team is willing to merge this PR)

[codecov-commenter]

Codecov Report

Attention: Patch coverage is 50.00000% with 2 lines in your changes missing coverage. Please review.

Project coverage is 74.45%. Comparing base (084b7e1) to head (afe9f86).
Report is 123 commits behind head on master.

Files with missing lines	Patch %	Lines
keras/src/trainers/trainer.py	50.00%	1 Missing and 1 partial ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##           master   #20315      +/-   ##
==========================================
- Coverage   78.83%   74.45%   -4.38%     
==========================================
  Files         512      512              
  Lines       48995    49062      +67     
  Branches     9021     9035      +14     
==========================================
- Hits        38624    36529    -2095     
- Misses       8509    10728    +2219     
+ Partials     1862     1805      -57     

Flag	Coverage Δ
keras	74.32% <25.00%> (-4.37%)	⬇️
keras-jax	62.26% <0.00%> (-0.04%)	⬇️
keras-numpy	?
keras-tensorflow	63.55% <25.00%> (-0.02%)	⬇️
keras-torch	62.25% <0.00%> (-0.04%)	⬇️

Flags with carried forward coverage won't be shown. Click here to find out more.

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

[fchollet]

Thanks for the PR! Yes, this seems like a reasonable change. Please take a look at the test failure: FAILED keras/src/trainers/trainer_test.py::TrainerDistributeTest::test_end_to_end_tf_distribute - RuntimeError: Random ops require a seed to be set when determinism is enabled. Please set a seed before running the op, e.g. by calling tf.random.set_seed(1).

[Frightera]

Hi @fchollet, I have no idea why jax backend test failed since I specially decorated the test for tensorflow backend:

keras/keras/src/trainers/trainer_test.py

Lines 1721 to 1741 in f24c863

	pytest.mark.skipif(
	backend.backend() != "tensorflow",
	reason="This test is only applicable to TensorFlow.",
	)
	@pytest.mark.requires_trainable_backend
	def test_jit_compile_with_tf_determinism(self):
	from tensorflow.python.framework.config import disable_op_determinism
	from tensorflow.python.framework.config import enable_op_determinism
	enable_op_determinism()
	model = ExampleModel(units=3)
	model.compile(
	optimizer=optimizers.SGD(),
	loss=losses.MeanSquaredError(),
	metrics=[metrics.MeanSquaredError()],
	)
	self.assertFalse(model.jit_compile)
	disable_op_determinism()

If I remove @pytest.mark.requires_trainable_backend, numpy backend fails on the same test. Any ideas?

[fchollet]

You forgot the @ for the decorator

[Frightera]

:-) it happens...