[udbhav-shrinet]

Summary

Fixes #443.

The flax and xreg optional dependency groups in pyproject.toml previously hardcoded jax[cuda], which breaks pip install timesfm[flax] or pip install timesfm[xreg] on CPU-only machines (laptops, CI runners, CPU cloud instances) where NVIDIA CUDA drivers are unavailable. Furthermore, jax[cuda] without a version qualifier is deprecated upstream in favor of jax[cuda12].

This PR splits each into:

• CPU-compatible default extras: flax and xreg using standard jax (defaults to CPU and works everywhere).
• Explicit GPU extras: flax-cuda and xreg-cuda using jax[cuda12].

Changes

• [project.optional-dependencies]:
  • flax: switches jax[cuda] -> jax
  • flax-cuda: added with jax[cuda12]
  • xreg: switches jax[cuda] -> jax
  • xreg-cuda: added with jax[cuda12]

Verification

• Verified pyproject.toml parses cleanly and extras are valid.
• CPU installations can now resolve flax and xreg without requiring CUDA drivers.

[google-cla]

Thanks for your pull request! It looks like this may be your first contribution to a Google open source project. Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA).

View this failed invocation of the CLA check for more information.

For the most up to date status, view the checks section at the bottom of the pull request.

[udbhav-shrinet]

CLA DONE !!

[udbhav-shrinet]

@googlebot I signed it.

[udbhav-shrinet]

@googlebot I have signed the CLA. Please recheck.

[udbhav-shrinet]

@googlebot I have signed the CLA. Please recheck.