ROCm/jax

Pull Requests Commits

charleshofer committed 345 days ago

500d314d

charleshofer committed 345 days ago

6a8b68ec

Add workflow disp event

charleshofer committed 345 days ago

a6c3a70c

Add test workflow for upstream runner

charleshofer committed 345 days ago

5a6a60d3

Merge pull request #219 from ROCm/ci-upstream-sync-104_1

github-actions[bot] committed 346 days ago

Verified 0962b960

Merge branch 'rocm-main' into ci-upstream-sync-104_1

charleshofer committed 346 days ago

ec2c3050

Merge pull request #25812 from Cjkkkk:segment_ids

Google-ML-Automation committed 346 days ago

7164c6ba

Update XLA dependency to use revision

Google-ML-Automation committed 346 days ago

49821e81

Add FFI example demonstrating the use of XLA's FFI state.

dfm committed 346 days ago

28afd252

Cjkkkk committed 346 days ago

8c4d6d69

Reject invalid None in jax.NamedSharding(spec=None).

pschuh committed 346 days ago

cb188a0c

[pallas_mgpu] For loops can have **non-ref** accumulators for carries.

cperivol committed 347 days ago

b48d15d7

Stop passing None into jax.NamedSharding in preparation for followup which bans passing None in (in favor of PartitionSpec()

pschuh committed 347 days ago

da97ee25

Reverts 3f737588fd49fcac36d6bd00cfdabae07df1fbb2

hawkinsp committed 347 days ago

58d8a97f

Merge pull request #26243 from jakevdp:einsum-asarray

Google-ML-Automation committed 347 days ago

57fa3721

Update XLA dependency to use revision

Google-ML-Automation committed 347 days ago

af84143e

Merge pull request #26174 from skye:cpu_configs

Google-ML-Automation committed 348 days ago

de48ce2a

[pallas/pallas_mgpu] Discharging run_scoped should not be discharging the intermediates

cperivol committed 348 days ago

b23f8f41

[pallas] Support DMA start partial discharge and run_scoped() does its own partial discharge.

cperivol committed 348 days ago

8649132d

Update XLA dependency to use revision

Google-ML-Automation committed 348 days ago

eb04fcbe

[Pallas TPU] Support jnp.take_along_axis for 32-bit vreg-sized vector.

bythew3i committed 348 days ago

ed952c8e

[Mosaic TPU] Support dynamic gather along axis 0 or 1 for 32-bit vreg-sized vector.

bythew3i committed 349 days ago

d8b92113

Move `ragged_all_to_all` test under appropriate test file

ghpvnist committed 349 days ago

c4e17632

Merge pull request #25766 from carlosgmartin:nn_initializers_variance_scaling_mode_fan_geo_avg

Google-ML-Automation committed 349 days ago

872e6c0e

Merge pull request #26249 from jakevdp:fix-sterling

Google-ML-Automation committed 349 days ago

a9f4dd71

Add mode='fan_geo_avg' to nn.initializers.variance_scaling.

carlosgmartin committed 349 days ago

96d3447e

Reverts bb951136e9b91a584bb422119ada76cc69c86024

emilyfertig committed 349 days ago

3b2410f7

Merge pull request #26195 from vfdev-5:use-tsan-numpy-in-tsan-ci

Google-ML-Automation committed 349 days ago

3f737588

Fix dtype issue in stirling approximation

jakevdp committed 349 days ago

216bd9a6

Merge pull request #26247 from hawkinsp:tsan2

Google-ML-Automation committed 349 days ago

0bf8462c

Older