ROCm/jax

Pull Requests Commits

Merge pull request #22836 from superbobry:maint-2

a-googler committed 1 year ago

0f103d33

Merge pull request #22870 from google:dependabot/github_actions/actions/upload-artifact-4.3.5

a-googler committed 1 year ago

af1a69ed

Bump actions/upload-artifact from 4.3.3 to 4.3.5

dependabot[bot] committed 1 year ago

Verified 6d7cf3fc

Merge pull request #22844 from jakevdp:xla-extension

a-googler committed 1 year ago

5f99c557

Better docs for jnp.fft.hfft and jnp.fft.ihfft

rajasekharporeddy committed 1 year ago

1acff9c7

Added a custom JVP rule for jax.nn.logsumexp

superbobry committed 1 year ago

0a48aca9

Move CostEstimate from pltu to pl

Google-ML-Automation committed 1 year ago

9762ac53

Nit. Fix missing backticks documentation for jnp.where.

hamelphi committed 1 year ago

ecf9f642

Merge pull request #22141 from dfm:update-cuda-call-example-to-ffi-call

a-googler committed 1 year ago

44a8c989

Merge pull request #22849 from jakevdp:dlpack-doc

a-googler committed 1 year ago

58927470

[pallas] Improve error and debugging messages with source locations

gnecula committed 1 year ago

252032a3

Port Eigenvalue Decompositions to XLA's FFI

Paweł Paruzel committed 1 year ago

b2a469b3

[pallas] Enable check for GPU lowering that tensor sizes are power of 2

gnecula committed 1 year ago

9b35b760

Update XLA dependency to use revision

Google-ML-Automation committed 1 year ago

0b87bf48

Reverts 80560663d3fab4c0c3f87d7c8e52fb9931526dbb

John Ryan committed 1 year ago

56ff247c

Merge pull request #22857 from mattjj:improve-while-loop-error

a-googler committed 1 year ago

83b5c7a0

Update XLA dependency to use revision

Google-ML-Automation committed 1 year ago

06c6a732

improve while_loop carry pytree/type mismatch errors

mattjj committed 1 year ago

bdcd358b

Improve documentation for jnp.from_dlpack

jakevdp committed 1 year ago

4d637c8f

Tighten the public API for jax.lib.xla_client & xla_extension

jakevdp committed 1 year ago

521c94c6

Don't call `api.clean_up` when there is no default backend.

yueshengys committed 1 year ago

09beb332

Fix lint in run_single_gpu.py

yueshengys committed 1 year ago

eb571c98

Update XLA dependency to use revision

Google-ML-Automation committed 1 year ago

51abbf90

Merge pull request #22828 from ROCm:ci_add_packages_dockerfile

a-googler committed 1 year ago

261cf520

Merge pull request #22784 from pearu:pearu/accuracy-tests-update

a-googler committed 1 year ago

e7fd424d

Update complex functions accuracy tests

pearu committed 1 year ago

780b10b4

Add `util.clear_all_caches` to `api.clear_backends` and let `api.clear_backends` be called before process terminates on JAX CPU. This could make the PjRt CPU client object to be successfully destroyed during Python garbage collection.

yueshengys committed 1 year ago

88c8bacd

Thread the mesh context manager to the place where we recover out_shardings back from GSPMDShardings. Before if you had a program like this:

yashk2810 committed 1 year ago

958234a9

[jax] Shard pallas_vmap_test

ezhulenev committed 1 year ago

ac52890e

Merge pull request #22840 from jakevdp:fix-backends

a-googler committed 1 year ago

6c79c104

Newer Older