ROCm/jax

Pull Requests Commits

Stop using xla/status.h, xla:status, and xla::Status now that xla::Status is just an alias for an absl::Status

klucke committed 1 year ago

ebdafea9

Merge pull request #21911 from gnecula:poly_solve_doc

a-googler committed 1 year ago

21e2319e

Rewrite `vector.multi_dim_reduction` with bf16 source/accumulator/output into

Google-ML-Automation committed 1 year ago

f86cd6de

Merge pull request #21904 from tilakrayal:patch-2

a-googler committed 1 year ago

039b7c12

Merge pull request #21901 from jjyyxx:patch-1

a-googler committed 1 year ago

2f22d3ab

Use ``unitialized_value`` for allocating outputs for interpreted Pallas kernels

superbobry committed 1 year ago

01f182e7

Batch `pxla.shard_args` calls triggered by `jax.device_put`

junwhanahn committed 1 year ago

cec796f5

[shape_poly] Add documentation for workaround with dimension parameters.

gnecula committed 1 year ago

b1a8c658

Added some docs to ``_hoist_consts_to_refs``

superbobry committed 1 year ago

550862f8

Use constant memory to pass in TMA descriptors to the kernel

apaszke committed 1 year ago

4ea73bf7

Fixed a ``ValueError`` when a Pallas GPU kernel closed over array constants

superbobry committed 1 year ago

f67f2e06

Rollback #21888, because it breaks multiple internal users

superbobry committed 1 year ago

4913fff9

rajasekharporeddy committed 1 year ago

b93da387

Fixing the naming conventions in lax_numpy.py

tilakrayal committed 1 year ago

Verified b461846c

Update XLA dependency to use revision

Google-ML-Automation committed 1 year ago

595a6208

Merge pull request #21897 from matsen:patch-1

a-googler committed 1 year ago

ae0127d6

Fix CUDNN requirement inconsistency in installation.md

jjyyxx committed 1 year ago

Verified cd23b2c8

Update XLA dependency to use revision

Google-ML-Automation committed 1 year ago

546a3a60

Merge pull request #21880 from gnecula:doc_poly

a-googler committed 1 year ago

993fa610

[shape_poly] Add documentation for shape polymorphism

gnecula committed 1 year ago

b58ff2ba

Fixing installation docs: needed quotes; markdown fix

matsen committed 1 year ago

Verified b585ecee

Further reduce the run time of `pmap_test`

Google-ML-Automation committed 1 year ago

50191671

Bump up tolerance for tests using `_testQdwh`.

bchetioui committed 1 year ago

4e748e1f

Merge pull request #21888 from jakevdp:lax-mul-bool

a-googler committed 1 year ago

193591b5

lax.mul: accept boolean inputs

jakevdp committed 1 year ago

4f7cd038

Update XLA dependency to use revision

Google-ML-Automation committed 1 year ago

895b4906

Merge pull request #21890 from hawkinsp:nightly

a-googler committed 1 year ago

99775043

Update nightly installation instructions to reflect new jax/jaxlib dependency structure.

hawkinsp committed 1 year ago

d5844b7c

Merge pull request #21856 from jakevdp:dtype-array

a-googler committed 1 year ago

e15e8b81

Merge all nightly wheels into `jax_nightly_releases.html` since we have dropped monolithic jaxlib CUDA wheels

yashk2810 committed 1 year ago

8f315c21

Newer Older