model.compile(jit_compile=True) for PyTorch has no intended effect (Keras 3.3)

[lseongjoo]

Description

When training a model in Keras 3.3 with the jit_compile=True option, a warning is displayed:

UserWarning: Model doesn't support `jit_compile=True`. Proceeding with `jit_compile=False`

PyTorch reports back its backends

assert torch.cuda.is_available()
torch.compiler.list_backends()
# ['cudagraphs', 'inductor', 'onnxrt', 'openxla', 'openxla_eval', 'tvm']

Investigation

Upon reviewing the source code, it appears that the jit_compile option is not implemented for PyTorch models. This feature might be planned but is not yet reflected in the code.

Steps to Reproduce

1. Define a model using PyTorch.
2. Compile the model with jit_compile=True.
3. Observe the warning during training.

Expected Behavior

The model should utilize jit_compile=True without warnings, or the documentation should clarify that this option is not supported for PyTorch models.

Environment

• Keras version: 3.3
• Backend: PyTorch 2.2

Additional Information

Please provide any updates on the planned implementation of jit_compile for PyTorch models or correct the documentation if this feature is not intended to be supported.

[sachinprasadhs]

Hi for the API's or Layers which does not support XLA then the model will not proceed with jit_compile=True, it will instead fallback to jit_compile=False mode.

[github-actions[bot]]

This issue is stale because it has been open for 14 days with no activity. It will be closed if no further activity occurs. Thank you.

[github-actions[bot]]

This issue was closed because it has been inactive for 28 days. Please reopen if you'd like to work on this further.

[google-ml-butler[bot]]

Are you satisfied with the resolution of your issue? Yes No