ROCm/jax

Pull Requests Commits

Update jax.host_id() and jax.host_count() to jax.process_index() and jax.process_count(), respectively.

danielsuo committed 23 days ago

46c02566

Merge pull request #34742 from mattjj:remove-zero-dot-from-primal-value

Google-ML-Automation committed 23 days ago

9c5d90dd

[Mosaic GPU] Fix handling of sliced/transposed types in `transformed_smem_ref_type`.

bchetioui committed 23 days ago

139f212f

Fix ValAccum to initialize Zero using .to_cotangent_aval()

yashk2810 committed 23 days ago

12462e05

[Mosaic GPU][NFC] Linter + type checker fixes in layout inference code.

bchetioui committed 23 days ago

457340ea

[Mosaic GPU][NFC] Delete dead code and clean up linter/pytype warnings in `layout_inference.py`.

bchetioui committed 23 days ago

fac19bbf

mattjj committed 23 days ago

39cc010b

[Mosaic GPU] Fix layout inference for `tcgen05_mma` to be compatible with user-specified tile sizes.

bchetioui committed 23 days ago

d751a3f7

[Mosaic TPU] Update some lowering rules to avoid using `tpu.repeat` since the op is going to be deprecated.

yueshengys committed 23 days ago

b1359ede

[ci] add pyrefly pre-commit hook

Jake VanderPlas committed 23 days ago

382ed1b7

Add the infer_params cache back that was removed during the lu.cache -> weakref_lru_cache move.

yashk2810 committed 23 days ago

b582e359

[Mosaic GPU] Fix bug in layout inference of WGMMA.

bchetioui committed 23 days ago

e786d22e

Clean up `LoadedExecutable::GetReadyFuture()`

junwhanahn committed 23 days ago

f0eddf93

[indexing] avoid printing full array repr in indexing-related error messages.

Jake VanderPlas committed 23 days ago

17b1bfe7

Fix test behavior.

yueshengys committed 23 days ago

54add9ec

Removed a few unused functions

superbobry committed 23 days ago

b58b083e

Merge pull request #34710 from jakevdp:ref-indexing

Google-ML-Automation committed 23 days ago

ab376494

Reverts f3ac8813699fba956b85702813795e36867b0833

Google-ML-Automation committed 23 days ago

fc1b3213

[refs] fix ref indexing corner case

jakevdp committed 23 days ago

c2899d0c

[mosaic] Avoid using the deprecated `output_memory_colors` in `CustomCallBackendConfig`

superbobry committed 23 days ago

f3ac8813

Update XLA dependency to use revision http://github.com/openxla/xla/commit/d2cdbae7b31f4a3b8c1144162f96e51b86b3f483

Google-ML-Automation committed 23 days ago

9489ca75

[Mosaic TPU] Support arbitrary offset during mask generation for 16-bit data in `SingleRowVRegBounds` by leveraging the 16-bit tpu.iota op.

yueshengys committed 23 days ago

480b5481

[Pallas/interpreter] Fix incorrect nesting of acquisitions of locks/condition variables.

Google-ML-Automation committed 23 days ago

0b8b8cea

Disable the `contracting2` test case from `test_dot_general`, following cl/859306670.

ZixuanJiang committed 24 days ago

875a7d6d

[Pallas/interpreter] Add an internal consistency check at the end of executing interpreted GPU kernels.

Google-ML-Automation committed 24 days ago

dd96fee0

Merge pull request #34713 from mattjj:empty-ref-p

Google-ML-Automation committed 24 days ago

d3753d4b

[mutable-arrays] add empty_ref primitive

mattjj committed 24 days ago

9dad8d5a

[Mosaic:TPU] Improve tpu.bitcast_vreg folding

tlongeri committed 24 days ago

88a07221

Move check_vma check inside pvary traceable instead of it's abstract_eval rule making it a complete no-op. Before pvary as a primitive was still in the jaxpr so it's transpose rule could be triggered.

yashk2810 committed 24 days ago

f037af0c

Allow unreduced/reduced to work when check_vma=False.

yashk2810 committed 24 days ago

11fe9d20

Newer Older