Pedro Goncalves Mokarzel

XPK install broken due to dependency

3

I am following the steps in https://github.com/AI-Hypercomputer/xpk?tab=readme-ov-file#installation after following the [pre-requisites](https://github.com/AI-Hypercomputer/xpk?tab=readme-ov-file#installation). Running `pip install xpk` will get the error: ``` ERROR: pip's dependency resolver does not currently take into account...

Enable lazy tensor loading for sharded tensors

11

Currently in PyTorchXLA when a tensor is initialized through sharding, it is loaded into its related devices immediately. From a logic point, `mark_sharding` is acting similarly to how calling `.to('xla')`...

Currently [`get_op_sharding`](https://github.com/pytorch/xla/blob/r2.7/torch_xla/distributed/spmd/xla_sharding.py#L116) generates an `xla::OpSharding`. With the new abstraction of `torch_xla::OpSharding`, we will want to use it instead

This is primarily for the sake of documentation and consistency.

Once we have refactored mark_sharding to utilize torch_xla::OpSharding, we will leverage it to implement Local SPMD. Through it we will store the correct global device association, and pass it to...

EDIT: Rather than creating an new RFC, I have decided to expand this GitHub issue with more information on achieving Local SPMD ## Context Previous work has been done to...

Have scripts/update_deps.py update bazel pin

1

Currently `scripts/update_deps.py` does not update bazel version. We should consider changing this to make sure pin updates done weekly also update bazel.

WORKSPACE is being migrated to Bzlmod (see the [migration bazel article](https://bazel.build/external/migration)). The WORKSPACE file is already disabled in Bazel 8 (late 2024) and will be removed in Bazel 9 (late...

As part of the "xla::OpSharding", I found two instances where it was actually being abstracted within PytorchXLA: - tensor_common.h: torch_xla::ShardingSpec - tensor.h/cpp: ShardingSpec Both instances do basically the same thing,...

Simplify device count external API calls

4

Currently there are many external APIs related getting the number of devices associate with PyTorch XLA. Those that I could find were: - "global_runtime_device_count": returns the total number of devices...