fix(tests): skip test_model_loading when the [flax] extra is missing

[LngelKyo]

tests/test_model_loading.py imports the flax backend at module top level:

from timesfm.timesfm_2p5.timesfm_2p5_flax import TimesFM_2p5_200M_flax

On a torch-only install (pip install ".[torch]"), flax and einshape are
absent (both are optional extras), so pytest fails at collection time and the
59 unrelated tests are lost with it (raised by sylvesterkaczmarek in review: three
of the four tests in this file are PyTorch-only, so a module-level skip also
removed the Torch model-loading and torch.compile coverage).

Gate only the one flax test on its optional dependencies. The flax import and both
importorskip calls now live inside test_flax_model_init_kwargs.

a. minimal install (no flax extra):

PYTHONPATH=src python -m pytest tests/test_model_loading.py -v
# test_torch_load_checkpoint_and_from_pretrained_local PASSED
# test_torch_compile_wraps_forward PASSED
# test_torch_no_compile_leaves_forward_unchanged PASSED
# test_flax_model_init_kwargs SKIPPED — "requires the [flax] extra"
# 3 passed, 1 skipped

b. with the [flax] extra installed:

PYTHONPATH=src python -m pytest tests/test_model_loading.py -v
# 4 passed

Full suite (PYTHONPATH=src python -m pytest tests -q): 63 passed. The skip
reason names the missing extra, so a user who actually wants the flax test gets
an actionable message. Related: #480.

[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.

[sylvesterkaczmarek]

The module-level pytest.importorskip skips this entire file when Flax/einshape is absent, but three of the four tests are PyTorch-only. On the minimal install this PR is targeting, that removes all Torch model-loading and torch.compile coverage too. Could the skip be scoped to test_flax_model_init_kwargs (or the imports split) so the Torch tests still run?

[LngelKyo]

Right, the module-level skip took the three torch tests with it. The flax import and both importorskip calls now live inside test_flax_model_init_kwargs, so a minimal install runs the torch model-loading and torch.compile tests (3 passed, 1 skipped) and a [flax] install runs all four. Verified both ways; commands and output in the PR description.

[jtjsvzd9md-alt]

ر #488