6 Pull requests merged by 2 people

• Set scoped vmem for paged attention

  #8988 merged Apr 17, 2025

• Fix merge conflict

  #8991 merged Apr 17, 2025

• Add an option for JittableModule to dedup parameters.

  #8965 merged Apr 16, 2025

• Add a helper class to handle mesh and sharding

  #8967 merged Apr 16, 2025

• Fix splash attention test

  #8978 merged Apr 16, 2025

• Fix call_jax hashing

  #8981 merged Apr 16, 2025

5 Pull requests opened by 3 people

• composibility of assume_pure and call_jax

  #8989 opened Apr 16, 2025

• typo fix

  #8990 opened Apr 16, 2025

• handle requires_grad in torchax

  #8992 opened Apr 17, 2025

• Set scoped vmem for paged attention (#8988)

  #8994 opened Apr 17, 2025

• cherry pick a37a5852808861a6b274cfbd1a98653fd1bc6a9a

  #8995 opened Apr 17, 2025

10 Issues closed by 3 people

• Using `StableHLOCompositeBuilder` in nested way.

  #6978 closed Apr 17, 2025

• [torchbench] The official benchmark for performance and accuracy check

  #7040 closed Apr 17, 2025

• model.to(xla_device) increases the number of named_parameters

  #7042 closed Apr 17, 2025

• The behavior of `torch.einsum` significantly differs between TPU and other devices.

  #7050 closed Apr 17, 2025

• multi_tensor_sgd triggers extra xla execution

  #7051 closed Apr 17, 2025

• Export nn.Module.forward with kwargs to StableHLO

  #7056 closed Apr 17, 2025

• (INACTIVE) TPU Freezing on loss.backward() on same epoch.

  #7101 closed Apr 17, 2025

• Problem with mesh shape in HybridMesh on TPU

  #7102 closed Apr 17, 2025

• Ragged paged attention test fail in v5e

  #8987 closed Apr 17, 2025

• [torchax] jit compile the model constructor

  #8635 closed Apr 16, 2025

1 Issue opened by 1 person

• Is there a way to attach metadata to a layer in a way that is included in the StableHLO export?

  #8993 opened Apr 17, 2025

25 Unresolved conversations

Sometimes conversations happen on old items that aren’t yet closed. Here is a list of all the Issues and Pull Requests with unresolved conversations.

• `call_jax` doesn't take jax config into hashing.

  #8963 commented on Apr 17, 2025 • 0 new comments

• torch.linalg.lstsq issues on GPU/TPU

  #8953 commented on Apr 17, 2025 • 0 new comments

• Support cost analysis in PyTorch/XLA

  #8750 commented on Apr 17, 2025 • 0 new comments

• 2.7 backport PR request list

  #8829 commented on Apr 17, 2025 • 0 new comments

• MegaScale discovery is ran twice again

  #8954 commented on Apr 17, 2025 • 0 new comments

• Splash attention test fail randomly only in github CI

  #8971 commented on Apr 17, 2025 • 0 new comments

• [call_jax] Bridge the torch_xla and JAX mesh

  #8972 commented on Apr 17, 2025 • 0 new comments

• `torch_xla::runtime::PjRtComputationClient::ExecuteReplicated()` crash after RuntimeError

  #8973 commented on Apr 17, 2025 • 0 new comments

• [assume_pure] test_assume_pure_complex_module torchax output on TPU differs from torch_xla

  #8974 commented on Apr 17, 2025 • 0 new comments

• [assume_pure] test_assume_pure_einsum_grads torchax output on TPU differs from torch_xla

  #8975 commented on Apr 17, 2025 • 0 new comments

• [torchax] `jax_tensor.copy_(cpu_tensor)` fails

  #8980 commented on Apr 17, 2025 • 0 new comments

• [torchax] requires_grad in some tensor constructors not propagated

  #8983 commented on Apr 17, 2025 • 0 new comments

• TPU test flake: test_diagonal_write_transposed_r3 (__main__.TestAtenXlaTensor)

  #8985 commented on Apr 17, 2025 • 0 new comments

• When trying to run this code with connection to tpu in google colab i had this error: AssertionError: 4 results for replica 0

  #8986 commented on Apr 17, 2025 • 0 new comments

• [RFC] PyTorch/XLA eager mode as default

  #7253 commented on Apr 17, 2025 • 0 new comments

• Cannot Import _XLAC

  #7070 commented on Apr 17, 2025 • 0 new comments

• Experiencing slow recompilation when manually building XLA

  #7057 commented on Apr 17, 2025 • 0 new comments

• torchdynamo + XLA crash

  #7053 commented on Apr 17, 2025 • 0 new comments

• Spmd whether expert parallelism is supported？

  #7049 commented on Apr 17, 2025 • 0 new comments

• Migrate PyTorch/XLA's gradient checkpointing to upstream one

  #7024 commented on Apr 17, 2025 • 0 new comments

• [Deprecation Tracking] API deprecation timeline summary

  #8915 commented on Apr 17, 2025 • 0 new comments

• Alternative to torch.select_mask

  #8968 commented on Apr 17, 2025 • 0 new comments

• Update test_export_fx_passes.py

  #8933 commented on Apr 16, 2025 • 0 new comments

• Default explicit donation for step barriers

  #8982 commented on Apr 17, 2025 • 0 new comments

• scan-based GRU falls back to nn.GRU when `bidirectional` is true.

  #8984 commented on Apr 17, 2025 • 0 new comments