Can only set jit compile WA for CPU device?

[guizili0]

System information.

TensorFlow version (you are using): TF2.12
Are you willing to contribute it (Yes/No) : Yes

Describe the feature and the current behavior/state.

https://github.com/keras-team/keras/blob/master/keras/layers/convolutional/base_conv.py#L285
use jit_compile as WA for CPU device, but pluggable device did not support XLA, that cause pluggable device cannot use keras group convolution

Will this change the current api? How?
No, just want to limit internal WA.

Who will benefit from this feature?
pluggable device user.

Contributing

• Do you want to contribute a PR? (yes/no): yes
• If yes, please read this page for instructions
• Briefly describe your candidate solution(if contributing):
  Add limit in https://github.com/keras-team/keras/blob/master/keras/layers/convolutional/base_conv.py#L285 that only for CPU.
  Just working on how to limit this in CPU.

[qlzh727]

What's the "WA" in this bug?

[guizili0]

What's the "WA" in this bug?

Here is the comments I got from the code: https://github.com/keras-team/keras/blob/master/keras/layers/convolutional/base_conv.py#L272

# TODO(b/213173659): remove this when grouped convolutions are fully
# supported on the CPU for compiled functions. For now, we need this as a
# workaround for CPU support.

[qlzh727]

@kenfranko for XLA as well.

[kenfranko]

Unclear to me how this is related to XLA. Is there something in particular that is the issue?

[guizili0]

Unclear to me how this is related to XLA. Is there something in particular that is the issue?

For this code, you add a WA(_jit_compiled_convolution_op) to support CPU groups convolution, but in TF plugin device, XLA still not support, so pluggable device would has issue here.
So what I want is just limit to use this WA when work on CPU. thanks.

[sachinprasadhs]

Thank you for reporting the issue.
We're now prioritizing Keras 3 development, a multi-backend framework that offers significant advantages by supporting JAX, TensorFlow, PyTorch and OpenVINO.
This means we won't be adding new features to TF-Keras (Keras 2).

TensorFlow users can benefit directly from Keras 3, which is included by default starting with version 2.16.1.
We encourage you to migrate to Keras 3 and experience its enhanced capabilities. If you feel this specific feature is crucial for the Keras framework itself, please open a new issue at: https://github.com/keras-team/keras/issues

For more details and usage guide about Keras 3, refer: https://keras.io/getting_started/about/