pytorch/xla

Pull Requests Commits

Refactor Computation of jax into creation and call (#8789)

qihqi committed 281 days ago

Verified 53b3ab81

Remove pages_per_seq padding constrains and add MHA support (#8816)

bythew3i committed 281 days ago

Verified af16c210

Fix the ragged kernel bug (#8806)

vanbasten23 committed 282 days ago

Verified 6804c65f

Document xla override (#8802)

pgmoka committed 283 days ago

Verified 0407d6d1

Leverage placeholder tensors for GA API (#8808)

rpsilva-aws committed 283 days ago

Verified 5b7462c6

Manually register einsum on xla (#8801)

pgmoka committed 285 days ago

Verified 1db6c861

Misc fixes (#8804)

qihqi committed 286 days ago

Verified 7aea0585

Re-enable triton tests. (#8803)

ysiraichi committed 286 days ago

Verified 2c70a1c2

Lower `_conj_copy` and `alias` operation. (#8686)

ysiraichi committed 287 days ago

Verified 6c53a1ef

Fix ragged paged attention v2 to resolve the recompilation issue (#8797)

yaochengji committed 287 days ago

Verified d06a9c97

Pin update to 20250303 (#8788)

tengyifei committed 287 days ago

Verified 9b61c1ad

Use placeholder tensor in scan (#8785)

tengyifei committed 287 days ago

Verified a9279287

Use .backward() with in-place grad mutations for the GA API (#8768)

rpsilva-aws committed 287 days ago

Verified 4540d817

Integrate ragged paged attention v2 (#8791)

bythew3i committed 287 days ago

Verified 5644f44a

Minimal support for calling JAX from PyTorch/XLA (#8781)

tengyifei committed 288 days ago

Verified 17270e2f

Build torch_xla wheel in build script (#8782)

zpcore committed 289 days ago

Verified 00fac787

write _shard_map; refactor flash attention to support 5d inputs. (#8730)

qihqi committed 289 days ago

Verified cee0820e

Introduce a GRU module implemented with scan (#8777)

tengyifei committed 289 days ago

Verified 0017ff16

Fix build on Python 3.9 (#8759)

tengyifei committed 292 days ago

Verified c76e9497

Misc changes to make torchax runnable on GPU. (#8756)

qihqi committed 292 days ago

Verified 60001c74

[scan] Make sure inputs into `fn` are not `device_data` IR nodes (#8769)

tengyifei committed 292 days ago

Verified 2feb0ac7

Add sm_scale in ragged attention kernel (#8771)

yaochengji committed 292 days ago

Verified 2675e689

Revert "Change num_seqs type from int to torch.Tensor" (#8767)

yaochengji committed 293 days ago

Verified 2e4f073a

Enable default buffer donation for gradient accumulation (#8758)

rpsilva-aws committed 293 days ago

Verified 6712eb91

Change num_seqs type from int to torch.Tensor (#8736)

vanbasten23 committed 293 days ago

Verified b18a65fa

Lower `as_strided_copy` use fast path with `slice` (#8734)

zpcore committed 294 days ago

Verified 1ab8216a

Add start_trace and stop_trace API in profiler (#8743)

lsy323 committed 294 days ago

Verified b4ba17b5

Add --network=host to TPU docker build command (#8735)

bhavya01 committed 294 days ago

Verified 71d0ce65

Increase tolerance in test_scan_xla_patched_linear (#8749)

tengyifei committed 294 days ago

Verified ac430abf

Extend buffer donation aliasing APIs (#8721)

rpsilva-aws committed 294 days ago

Verified 07976424

Newer Older