ROCm/jax

Pull Requests Commits

Initial commit for `jax.experimental.compute_on` API.

yashk2810 committed 1 year ago

2d6d408b

Merge pull request #21289 from jakevdp:stacklevel

a-googler committed 1 year ago

7a3fc711

jax.experimental.export: fix stacklevel for warning

jakevdp committed 1 year ago

9ad6729b

Merge pull request #21283 from dfm:gh21279

a-googler committed 1 year ago

56301d47

[pallas:Mosaic GPU] Configurable smem scratch and a small bug fix in Mosaic GPU

Google-ML-Automation committed 1 year ago

81525668

Add RegularGridInterpolator to generated API docs

dfm committed 1 year ago

be069547

Merge pull request #21281 from jaro-sevcik:enable-host-offloading-scan-tests-gpu

a-googler committed 1 year ago

e93f36aa

Enable remat-scan offloading test

jaro-sevcik committed 1 year ago

7be7c1ea

[Mosaic GPU] reduce_sum does an intra-warp reduction before communicating with the other warps

Google-ML-Automation committed 1 year ago

f87be35b

Use absl.flags.FlagHolder to defined --mosaic_gpu_debug

superbobry committed 1 year ago

210f8bbf

Internal BUILD file change

Google-ML-Automation committed 1 year ago

c4559115

Added a slow (but working!) implementation of layer norm in Pallas via Mosaic GPU

superbobry committed 1 year ago

527aef3a

Merge pull request #21261 from superbobry:mypy-ruff

a-googler committed 1 year ago

5e2710c2

Merge pull request #21268 from jakevdp:register-dataclass

a-googler committed 1 year ago

1829a667

[Mosaic GPU] Add a WGSplatLayout that trivially supports reshape and broadcast.

Google-ML-Automation committed 1 year ago

0e924339

Update XLA dependency to use revision

Google-ML-Automation committed 1 year ago

efa420b2

Merge pull request #21181 from 8bitmp3:add-distributed-data-loading-doc

a-googler committed 1 year ago

8493fc97

Document jax.tree_util.register_dataclass

jakevdp committed 1 year ago

defb53f4

Add Distributed data loading doc

8bitmp3 committed 1 year ago

15a3debe

Merge pull request #21267 from jakevdp:quantile-warning

a-googler committed 1 year ago

bc19f7f9

Bumped mypy to 1.10.0 and ruff to 0.4.4

superbobry committed 1 year ago

c3bc88d5

jnp.quantile & friends: properly deprecate interpolation

jakevdp committed 1 year ago

bfbde5ef

Merge pull request #21244 from google:add-tpu-core-count

a-googler committed 1 year ago

9dd98dcd

Switch the order of sharding and memory kind custom call

yashk2810 committed 1 year ago

74ffed90

Clarified the type of the inputs to callback APIs

superbobry committed 1 year ago

01194bd2

[Mosaic GPU] Enable pointwise operation with scalar values.

Google-ML-Automation committed 1 year ago

380503b0

Add a sharding test for concatenate with an internal operand that requires

bixia1 committed 1 year ago

2ec39941

Add jax.make_array_from_process_local_data to create a distributed tensor from host data and supporting scaffolding in sharding to be able to figure out dimensions of host data required.

marksandler2 committed 1 year ago

8f045caf

Use `util.cache` instead of `lru_cache` for `create_mesh_pspec_sharding`

junwhanahn committed 1 year ago

cd41b4fd

Update XLA dependency to use revision

Google-ML-Automation committed 1 year ago

6fe313c6

Newer Older