ROCm/jax

Pull Requests Commits

Regenerate python stubs for nanobind 2.11.0.

hawkinsp committed 119 days ago

6fb3f3da

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

Google-ML-Automation committed 119 days ago

3a93274b

Merge pull request #34462 from alekstheod:add_support_of_building_jax_under_umbrella_workspace

Google-ML-Automation committed 119 days ago

4a5ad80e

Adds `jax.lax.scaled_dot` for scaled dot products.

loislo committed 120 days ago

50b394dc

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

Google-ML-Automation committed 120 days ago

91871a92

[Mosaic GPU][NFC] Clean up the lowering of `slice_smem`.

bchetioui committed 120 days ago

01d76c43

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

danielsuo committed 120 days ago

46c02566

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

Google-ML-Automation committed 120 days ago

9c5d90dd

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

bchetioui committed 120 days ago

139f212f

Fix ValAccum to initialize Zero using .to_cotangent_aval()

yashk2810 committed 120 days ago

12462e05

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

bchetioui committed 120 days ago

457340ea

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

bchetioui committed 120 days ago

fac19bbf

mattjj committed 120 days ago

39cc010b

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

bchetioui committed 120 days ago

d751a3f7

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

yueshengys committed 120 days ago

b1359ede

[ci] add pyrefly pre-commit hook

Jake VanderPlas committed 120 days ago

382ed1b7

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

yashk2810 committed 120 days ago

b582e359

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

bchetioui committed 120 days ago

e786d22e

Clean up `LoadedExecutable::GetReadyFuture()`

junwhanahn committed 121 days ago

f0eddf93

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

Jake VanderPlas committed 121 days ago

17b1bfe7

Fix test behavior.

yueshengys committed 121 days ago

54add9ec

Removed a few unused functions

superbobry committed 121 days ago

b58b083e

Merge pull request #34710 from jakevdp:ref-indexing

Google-ML-Automation committed 121 days ago

ab376494

Reverts f3ac8813699fba956b85702813795e36867b0833

Google-ML-Automation committed 121 days ago

fc1b3213

[refs] fix ref indexing corner case

jakevdp committed 121 days ago

c2899d0c

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

superbobry committed 121 days ago

f3ac8813

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

Google-ML-Automation committed 121 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 121 days ago

480b5481

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

Google-ML-Automation committed 121 days ago

0b8b8cea

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

ZixuanJiang committed 121 days ago

875a7d6d

Newer Older