ROCm/jax

Pull Requests Commits

charleshofer committed 1 year ago

500d314d

charleshofer committed 1 year ago

6a8b68ec

Add workflow disp event

charleshofer committed 1 year ago

a6c3a70c

Add test workflow for upstream runner

charleshofer committed 1 year ago

5a6a60d3

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

github-actions[bot] committed 1 year ago

Verified 0962b960

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

charleshofer committed 1 year ago

ec2c3050

Merge pull request #25812 from Cjkkkk:segment_ids

Google-ML-Automation committed 1 year ago

7164c6ba

Update XLA dependency to use revision

Google-ML-Automation committed 1 year ago

49821e81

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

dfm committed 1 year ago

28afd252

Cjkkkk committed 1 year ago

8c4d6d69

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

pschuh committed 1 year ago

cb188a0c

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

cperivol committed 1 year ago

b48d15d7

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

pschuh committed 1 year ago

da97ee25

Reverts 3f737588fd49fcac36d6bd00cfdabae07df1fbb2

hawkinsp committed 1 year ago

58d8a97f

Merge pull request #26243 from jakevdp:einsum-asarray

Google-ML-Automation committed 1 year ago

57fa3721

Update XLA dependency to use revision

Google-ML-Automation committed 1 year ago

af84143e

Merge pull request #26174 from skye:cpu_configs

Google-ML-Automation committed 1 year ago

de48ce2a

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

cperivol committed 1 year ago

b23f8f41

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

cperivol committed 1 year ago

8649132d

Update XLA dependency to use revision

Google-ML-Automation committed 1 year ago

eb04fcbe

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

bythew3i committed 1 year ago

ed952c8e

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

bythew3i committed 1 year ago

d8b92113

Move `ragged_all_to_all` test under appropriate test file

ghpvnist committed 1 year ago

c4e17632

Merge pull request #25766 from carlosgmartin:nn_initializers_variance_scaling_mode_fan_geo_avg

Google-ML-Automation committed 1 year ago

872e6c0e

Merge pull request #26249 from jakevdp:fix-sterling

Google-ML-Automation committed 1 year ago

a9f4dd71

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

carlosgmartin committed 1 year ago

96d3447e

Reverts bb951136e9b91a584bb422119ada76cc69c86024

emilyfertig committed 1 year ago

3b2410f7

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

Google-ML-Automation committed 1 year ago

3f737588

Fix dtype issue in stirling approximation

jakevdp committed 1 year ago

216bd9a6

Merge pull request #26247 from hawkinsp:tsan2

Google-ML-Automation committed 1 year ago

0bf8462c

Older