pytorch/xla

Pull Requests Commits

yaochengji committed 1 year ago

82d089b9

seperate tests in another file

yaochengji committed 1 year ago

338db2d3

Run the MP tests on TPU

tengyifei committed 1 year ago

ff5e6603

Update jax to 0.5.4

tengyifei committed 1 year ago

61f35290

bump libtpu to 0.0.12 nightly

tengyifei committed 1 year ago

e3302034

update libtpu date to 20250320

yaochengji committed 1 year ago

b84f5c12

yaochengji committed 1 year ago

676f5b8d

yaochengji committed 1 year ago

5ffbd7a2

yaochengji committed 1 year ago

ed1f9f7e

yaochengji committed 1 year ago

b3d9739e

[ragged-paged-attn] Use hidden states in kv cache and support any num_kv_head (#8851)

bythew3i committed 1 year ago

Verified 4190fc0e

fix dependency (#8852)

zpcore committed 1 year ago

Verified c7d0b1ef

fix existing bazel link removal (#8848)

zpcore committed 1 year ago

Verified 18cdb17a

Make cxx_abi enable by default (#8844)

zpcore committed 1 year ago

Verified 8f314a5a

install jax dependency by default to unblock pytorch ci (#8845)

zpcore committed 1 year ago

Verified 4c5cc420

Add information about codegen in codegen/xla_native_functions.yaml (#8817)

pgmoka committed 1 year ago

Verified 76b0ce50

Add info related to codegen (#8834)

pgmoka committed 1 year ago

Verified 5dd50099

Expose mark_sharding_with_gradients as a public API. (#8826)

iwknow committed 1 year ago

Verified 5caaaee8

Fix for triton pin update. (#8825)

ysiraichi committed 1 year ago

Verified 818c92df

Add build triggers for 2.7-rc1 (#8828)

zpcore committed 1 year ago

Verified fed01200

Revamp SPMD guide (#8807)

tengyifei committed 1 year ago

Verified 5beb6511

Adapt Neuron data type tests (#8835)

rpsilva-aws committed 1 year ago

Verified bd46c0ea

Remove xla_mark_sharding_dynamo_custom_op (#8823)

tengyifei committed 1 year ago

Verified 27d455aa

Support unused arguments in xb.call_jax (#8830)

tengyifei committed 1 year ago

Verified b3f38757

[torchax] Fix functional.max_pool (#8814)

dvhg committed 1 year ago

Verified 0aac10e2

Proper test coverage for Neuron CollectivePermute (#8818)

rpsilva-aws committed 1 year ago

Verified acc4c739

Fix jax_env_context (#8819)

bhavya01 committed 1 year ago

Verified 2e04cfb4

Update MarkShardingFunction to be AOTAutograd traceable (#8821)

tengyifei committed 1 year ago

Verified c969276e

[Compilation Cache] Add API to query number of cached computation graphs (#8822)

lsy323 committed 1 year ago

Verified bdf30f66

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

qihqi committed 1 year ago

Verified 53b3ab81

Older