[TPU] Fix tpu torch compile error

[Chenyaaang]

Fix TPU torch compile error after #26113

This pr includes:

1. Set compilation backend to openxla on TPU platform
2. Make sure TPU is using forward_tpu when dispatching in custom ops
3. Bypass some backend checks which require either eager or inductor for non-tpu platforms.

[gemini-code-assist]

Code Review

I've reviewed the changes and they look good for fixing the TPU compilation issue. The logic to enable custom ops and set the openxla backend for TPU is correct. I have one suggestion to improve the user experience by adding a log message when the backend is overridden, which is consistent with other parts of the code.

[ProExpertProg]

Thanks for catching this, sorry we didn't catch this during review. The reason this got broken is because the logic is complicated, could we improve it?

[ProExpertProg]

Can you make this a property of current_platform? Perhaps current_platform.default_dynamo_backend?

[ProExpertProg]

Let's just try to move this after the platform-specific update, and TPU can set custom ops the way it wants to.

[Chenyaaang]

yes we can move the self.compilation_config.custom_ops logic after current_platform.check_and_update_config(self), and the new statement will be self.compilation_config.backend not in ["eager", "openxla"], does it look better to you?

[ProExpertProg]

I actually think that tpu platform should just do:

        if (
            self.compilation_config.custom_ops.count("none")
            + self.compilation_config.custom_ops.count("all")
            == 0
        ):
            self.compilation_config.custom_ops.append("all")

And then the common logic won't even be hit because custom_ops already contains "all" (once the logic is moved after).

[Chenyaaang]

After some investigations, I prefer not to move current_platform.check_and_update_config(self) up to before this part. I'm afraid the logic inside each platform's check_and_update_config might be related to some of its previous inits, and since my change is only related to TPU, I don't want to add more potential risks to other platforms.

[ProExpertProg]

Original PR was reverted, we can include these fixes in the unrevert

[Chenyaaang]

Original PR was reverted, we can include these fixes in the unrevert

SG, I'll go ahead and close this pr. Can you also cc me in the new pr? Thanks!