* Make python Api to respect the virtual device when SPMD is enabled

* fix typo

* Also dump output sharding on HLO file

* only dump output sharding if dump format is HLO

* add test

* fix typo

* Make all-reduce a no-op when world size is 1

* Fix torch.distributed test

* fix amp dtype setting for GPU.

* fix ut

* fix lint.

* minor.

* Add python test for SPMD+Runtime Python API

* replace test name

* Update test_xla_spmd_python_api_interaction.py

…5352)

* Check the actual device instead of query env var for virtual device

* revert unneeded change

* minor changes

…5364)

…ro1.py (#5364)" (#5366)

This reverts commit 8ada333.

* tweak `atol` and `rtol`

*  Skip`DynamoTrainingBasicTest.test_resnet18` on TPU

* Add kokoro presubmit for stablehlo tests

)


* Skip `DynamoInferenceBasicTest.test_resnet18` on TPU

…5367)

* [BE] use self.assertEquals instead of str equality in test_zero1.py

* Use our own assertEqual

* Remove print statements

* Fix ReplicateShardedData for int type

* add test

Update dynamo.md to remove note about fallback ops since they're supported now

This reverts commit 7fb7dfe.

…erent_input_shape` on TPU (#5373)

* tweak `atol` and `rtol` for `test_simple_model_with_different_input_shape` on TPU

…mutability (#5382)

* [TEST ONLY] print statements for test_zero1.py to debug

* Try fix

* Rectify test_zero1.py to account for state_dict modification

* Fix lint

#5384)

* Add gpu doc for how to build PyTorch/XLA from source with GPU support.

* fix typo

* fix comments

* fix comments

* Add more support for in-place ops in dynamo bridge

Run linter

* Add check to explicitly sync self tensors

Remove debugging lines

Update unit tests to a model

* Clean up some code

Surround  in an if-statement

Update metrics for fallback related dynamo tests

Update cloned args logic

Revert "Update metrics for fallback related dynamo tests"

This reverts commit 3855f43.

* Update single_node flag back to False

Add dynamo test in TPU CI

Summary:
During the LLaMA2 experiements, I disovered that manually marking 1D tensors to be replicated can greatly save a lot of memory. Then I disocvered that explicitly replicated spec will get dropped after mark_step. That is caused by PrepareOutputShardingPropagation where it explicitly clear the sharding spec for replicated output. So, I went ahead and fix that.

Further, I did some experiements of propogating replicated output and that drop the requirements of manually replicating 1D tensors. Hence, I made this change.

I'm still not quite sure why, will follow up later.

Test Plan:
PJRT_DEVICE=TPU python test/spmd/test_xla_sharding.py

Update more places

Add torch_pin

Summary:
This change enables megacore_dense by default to allow asynchorous cc
ops especailly for GSPMD.

Test Plan:
CI

Co-authored-by: Jiewen Tan <jwtan@google.com>

* Add optiona to unbundle libtpu

* Add clarifying note

* Fix fsdp not freeing forzen full params

* add test

* formatting

* remove unnecessary env var in test

Co-authored-by: Liyang90 <liyanglu@google.com>

* Update project metadata and remove useless files

* Update README

* Add manylinux platform tag

* formatting

* Add resnet50-weight-only-quant colab notebook

* update notebook with llama blog link

Co-authored-by: Siyuan Liu <lsiyuan@google.com>