ROCm/jax

Pull Requests Commits

[XLA:Mosaic] Support trunc/ext op for 1D vector with any implicit dim.

bythew3i committed 1 year ago

16716170

jax.debug.callback now passes arguments as jax.Arrays

superbobry committed 1 year ago

6e23c14f

Update outdated comment Pmap_Cookbook

sagelywizard committed 1 year ago

Verified c980dc4e

quickstart tweaks (from https://github.com/google/jax/pull/20819)

mattjj committed 1 year ago

94e3a6ef

Add a common flag for the collectives implementations on cpu.

inailuig committed 1 year ago

60d4c43f

jax.debug.callback now requires a Callable[..., None]

superbobry committed 1 year ago

32922f61

[PJRT C API] Add stream extension to support DLPack and implement this extension in CUDA plugin.

Jieying Luo committed 1 year ago

b2375fa7

[jax2tf] Adjust tolerance for asinh test.

gnecula committed 1 year ago

cea36a04

Merge pull request #20825 from jakevdp:gammasgn

a-googler committed 1 year ago

c7517b83

Finalize deprecation of zero-dimensional inputs to jnp.nonzero

Jake VanderPlas committed 1 year ago

41fa67c2

Cache the `_check_sharding` check in device_put. If aval and sharding are the same, no need to check multiple times

yashk2810 committed 1 year ago

837f0bbf

Update XLA dependency to use revision

a-googler committed 1 year ago

8fec8a60

Add jax.scipy.special.gammasgn

jakevdp committed 1 year ago

568db105

Merge pull request #20823 from sagelywizard:patch-1

a-googler committed 1 year ago

4d6c4d63

Remove outdated section in cloud_tpu_colabs README

sagelywizard committed 1 year ago

Verified afc87785

[Jax/Triton] Skip benchmarking while autotuning for configs that cannot be launched.

marvin-kim committed 1 year ago

90e9e47a

Merge pull request #20816 from superbobry:int4

a-googler committed 1 year ago

de7d3b68

Make `core.Token` a non-trivial class which wraps a `jax.Array`. Currently, we use a singleton and empty `core.token` object everywhere. After the change, tokens could be created and threaded in and out of computations to build up dependency.

yueshengys committed 1 year ago

c2d43735

[Pallas TPU] Generalize while_loop lowering in Pallas -> Mosaic.

a-googler committed 1 year ago

9c9e805e

Merge pull request #20537 from jakevdp:update-tpu-readme

a-googler committed 1 year ago

6ca69f38

Skip `test_spmd_preserves_input_sharding_vmap_grad` unless `xla_extension_version >= 258`

yashk2810 committed 1 year ago

aaac1d8f

Fix bug in rank-deficient fix-up code: Do not zero out the corresponding column of u_out if a diagonal entry of r is exactly zero.

a-googler committed 1 year ago

51763d8b

Added int4 and uint4 to dtype-specific tests

superbobry committed 1 year ago

a13efc28

Document the fact that jax.clear_caches() doesn't affect the persistent cache.

a-googler committed 1 year ago

bb8cf34a

Increase sharding to avoid timeouts

apaszke committed 1 year ago

fa66e731

Initial commit for Mosaic GPU

apaszke committed 1 year ago

8e3f5b10

Update XLA dependency to use revision

a-googler committed 1 year ago

c4dea624

Merge pull request #20806 from jakevdp:tweak-docs

a-googler committed 1 year ago

2a65cd38

DOC: one last readthrough of the new 101 tutorials

jakevdp committed 1 year ago

48e8457b

Merge pull request #20768 from chaserileyroberts:chase/eval_jaxpr_in_extend

a-googler committed 1 year ago

1a8aae0e

Newer Older