Remove excessive warn message in maybe_get_jax as it creates too many log lines during training

[rajkthakur]

🐛 Bug

The maybe_get_jax() function in torch_xla/_internal/jax_workarounds.py merged in #9521 currently emits a warning message when JAX is not installed. While informative, this warning results in an excessive number of log lines during training workloads, cluttering the logs and making it difficult to spot genuinely important debug messages.

To Reproduce

Steps to reproduce the behavior:

1. Create Python Virtual Environment (python3 -m venv ptxla_28) on Ubuntu 22.04
2. pip install torch==2.8.0 torchvision; pip install torch_xla==2.8.0
3. Create small python script(let's call it trigger_warning.py)

import sys
sys.path.insert(0, 'ptxla_28/lib/python3.10/site-packages')
from torch_xla._internal.jax_workarounds import maybe_get_jax
maybe_get_jax() 

5. execute the script bash -c "source ptxla_28/bin/activate && python trigger_warning.py"
6. You should be able to see the warning message like below

WARNING:root:Defaulting to PJRT_DEVICE=CPU
WARNING:root:You are trying to use a feature that requires jax/pallas.You can install Jax/Pallas via pip install torch_xla[pallas]

Expected behavior

Remove or suppress this warning message, or limit it to display only once per process/session instead of for every invocation.

Environment

• Reproducible on XLA backend [CPU/TPU/CUDA]: CPU
• torch_xla version: 2.8.0
• Relevant Code:
  

  xla/torch_xla/_internal/jax_workarounds.py

  Line 61 in 0f56dec

  def maybe_get_jax():

Additional context

The current behavior results in thousands of lines of repeated warnings when running workloads that do not require JAX, negatively impacting developer experience. Reducing or removing this warning will significantly clean up logs for users running long or large-scale training jobs, improving usability without sacrificing relevant error reporting.

[subhashgahlot161278]

+1

[jeffhataws]

Mark this as performance issue since affects lazy mode training performance for multiple models.

[ysiraichi]

cc @qihqi

[qihqi]

Hi,

Sure, we can have an option to disable this warning via env variables and what not.

However, usually a function calls maybe_get_jax is because it actually need to use jax (say pallas, shard_as, call_jax, assume_pure); if you can figure out which function called maybe_get_jax and is unneeded then that will be better, we can remove call to maybe_get_jax from that.

If you actually used a fucntion that needs jax, but you dont have jax, then that function likely is not doing what it is suppose to do (say shard_as becoming no-op) and that is no good.

[rajkthakur]

Since JAX is an optional dependency, we should introduce a new env variable that applies to maybe_get_jax so that the function return None if the env is not set.

Additional Context:
We have observed upto 20% performance degradation just because this check is in hot path of some of our test cases. So if we can return reliably using the env variable it helps restore the perf drop.

[qihqi]

@rajkthakur sounds good. Feel free to submit a PR