Introduce Model Parallel infrastructure and DTensor support for Keras Variables in torch core.py

[buildwithsuhana]

This PR implements distributed variable initialization and tensor promotion for the PyTorch backend using torch.distributed.tensor (DTensor). It enables ModelParallel distribution strategies to correctly shard variables and manage mixed-tensor operations.

Key Changes:

• Distributed Variable Initialization: Implemented _initialize_distributed in keras/src/backend/torch/core.py. This ensures that when a ModelParallel distribution is active, Variable objects are initialized as DTensors according to the specified layout.
• DTensor Promotion Mode: Introduced KerasDTensorPromotionMode, a TorchDispatchMode that automatically promotes plain tensors to replicated DTensors when they interact with DTensors. This prevents "mixed tensor" errors during distributed execution.
• Automatic Promotion Activation: convert_to_tensor now automatically enters the global KerasDTensorPromotionMode when it encounters a DTensor, ensuring seamless integration with standard Keras ops.
• Updated _direct_assign to handle DTensor targets, ensuring that assigned values are correctly distributed before copying.
• Updated convert_to_numpy to handle DTensors by first redistributing them to a replicated state and then converting the local data to NumPy.

Description

Contributor Agreement

Please review our AI-Assisted Contribution Policy and check all boxes below before submitting your PR for review:

• I am a human, and not a bot.
• I will be responsible for responding to review comments in a timely manner.
• I will work with the maintainers to push this PR forward until submission.

Note: Failing to adhere to this agreement may result in your future PRs no longer being reviewed.

[gemini-code-assist]

Code Review

This pull request introduces support for PyTorch DTensors and model parallelism in the PyTorch backend, including automatic tensor promotion, distributed variable initialization, and direct assignment. The review feedback focuses on improving robustness against device mismatch errors (e.g., CPU vs. CUDA) during tensor promotion, variable initialization, and direct assignment. It also suggests skipping device transfers for DTensors in convert_to_tensor to prevent runtime errors, and using a context manager for KerasDTensorPromotionMode in tests to avoid leaking the dispatch mode.

[codecov-commenter]

Codecov Report

❌ Patch coverage is 7.92079% with 93 lines in your changes missing coverage. Please review.
✅ Project coverage is 22.03%. Comparing base (c934679) to head (7ab6713).

Files with missing lines	Patch %	Lines
keras/src/backend/torch/core.py	7.14%	78 Missing ⚠️
...s/src/trainers/data_adapters/py_dataset_adapter.py	15.38%	11 Missing ⚠️
keras/src/backend/torch/distribution_lib.py	0.00%	4 Missing ⚠️

❗ There is a different number of reports uploaded between BASE (c934679) and HEAD (7ab6713). Click for more details.

HEAD has 27 uploads less than BASE

Flag	BASE (c934679)	HEAD (7ab6713)
keras	10	1
keras-tpu	2	1
keras-jax	4	1
keras-numpy	1	0
keras-cpu	5	0
keras-openvino	1	0
keras-tensorflow	2	0
keras-torch	2	0
keras-gpu	3	0

Additional details and impacted files

@@             Coverage Diff             @@
##           master   #23402       +/-   ##
===========================================
- Coverage   85.03%   22.03%   -63.00%     
===========================================
  Files         468      468               
  Lines       71147    71244       +97     
  Branches    11792    11819       +27     
===========================================
- Hits        60499    15698    -44801     
- Misses       7634    54565    +46931     
+ Partials     3014      981     -2033     

Flag	Coverage Δ
keras	22.03% <7.92%> (-62.81%)	⬇️
keras-cpu	?
keras-gpu	?
keras-jax	22.03% <7.92%> (-36.63%)	⬇️
keras-numpy	?
keras-openvino	?
keras-tensorflow	?
keras-torch	?
keras-tpu	22.03% <7.92%> (-35.47%)	⬇️

Flags with carried forward coverage won't be shown. Click here to find out more.

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[hertschuh]

Ok. But that's a really weird place to enter this context manager.

Should it be open when we set the global distribution?

Should it be open always when we initialize keras with Torch as a backend?

[hertschuh]

Move this import to the top

[buildwithsuhana]

Actually, this was kept as a local import to prevent circular import dependency. The distribution library imports small utilities from this module (e.g. get_device,_parse_device_input). Importing keras.src.backend.torch.distribution_lib at module top would create a cycle during import

[hertschuh]

Move this import to the top.

[buildwithsuhana]

This would also cause a circular dependency hence we import TensorLayout locally to avoid circular imports between the high-level distribution API and backend-specific modules.

[hertschuh]

Isn't that covered by lines 394-395?

Under what other circumstances would this be needed? Also, the if isinstance(x, Variable) should definitely come first.

[buildwithsuhana]

I've updated the code to remove that redundant check

[hertschuh]

Add comment about the fact that this context is intentionally never exited.

[buildwithsuhana]

Added

[hertschuh]

Just for my understanding, why was the else added here?

[buildwithsuhana]

Removed the else block for consistency :)

[hertschuh]

Thanks!

[hertschuh]

So it's crashing on Torch now. Can you look into it?