ROCm/jax

Pull Requests Commits

[ROCm] Rollup of changes for rocm7 support

mrodden committed 1 year ago

af2dcef8

Merge pull request #30962 from jakevdp:default-dtype-bits-dep

Google-ML-Automation committed 1 year ago

297e2a1c

[Mosaic] Add strict_ordering attribute to tpu.enqueue_dma and tpu.wait_dma2

naummo committed 1 year ago

81cf6da6

Disable tests that will fail with an upcoming llvm integrate.

Google-ML-Automation committed 1 year ago

49f3ff73

Added version guard for `recoverable` argument.

mwhittaker committed 1 year ago

3d79e2be

Update error message in a few places where we say `operand shape` when we mean `operand type`.

yashk2810 committed 1 year ago

48c766cd

Merge pull request #30980 from gspschmid:gschmid/ref-pp-like-val

Google-ML-Automation committed 1 year ago

26f0d129

Merge pull request #30979 from superbobry:main

Google-ML-Automation committed 1 year ago

df7503ab

Removed misplaced docstring

superbobry committed 1 year ago

a0230390

Replaced unnecessary usages of `functools.reduce` with `math.prod` or `sum`

superbobry committed 1 year ago

c44138ed

[Pallas:MGPU] Expose the optimized SMEM/GMEM copy layout

apaszke committed 1 year ago

b3948bfe

[pallas] Removed another dtype bit width function in favor of `dtypes.bit_width`

superbobry committed 1 year ago

8f87c00c

[JAX] Remove more uses of canonicalize_dtype.

hawkinsp committed 1 year ago

ecfb658f

[Mosaic GPU][NFC] Use `pass_boilerplate.h` to simplify `GpuLaunchLoweringPass`.

bchetioui committed 1 year ago

b3ba19f8

[Mosaic GPU] Only pass in -hex to nvdisasm when we actually care about the control codes

apaszke committed 1 year ago

541e9274

[Pallas/Mosaic GPU] Update documentation to mention `MOSAIC_GPU_DUMP_LLVM`.

bchetioui committed 1 year ago

fc8b2b46

[ref] Respect short_dtypes and mesh_axis_types when pretty-printing AbstractRefs

gspschmid committed 1 year ago

9f53b973

Removed more dead code

superbobry committed 1 year ago

bc6266f4

[Mosaic GPU] Add `MOSAIC_GPU_DUMP_LLVM` as a short-hand for `MOSAIC_GPU_LLVM_DEBUG_ONLY=serialize-to-llvm`.

bchetioui committed 1 year ago

8b332705

[Mosaic GPU][NFC] Inline in/out shape definitions.

allanrenucci committed 1 year ago

335a05f5

Update XLA dependency to use revision http://github.com/openxla/xla/commit/953ad92148941c10281dd764b78ac192c91761bf

Google-ML-Automation committed 1 year ago

7c2da79a

[Mosaic:TPU] Get rid of RectangularVregBounds and clean up VectorLayout::tileDataBounds

tlongeri committed 1 year ago

290b0a38

Automated Code Change

Google-ML-Automation committed 1 year ago

93b464d3

Add trivial batching rule for tpu custom call

sharadmv committed 1 year ago

8f706931

Merge pull request #30968 from mattjj:cond-linearize

Google-ML-Automation committed 1 year ago

e00df525

Merge pull request #30957 from jakevdp:compilation-cache-dep

Google-ML-Automation committed 1 year ago

c7586a47

Use standard deprecation path for jax.experimental.compilation_cache

jakevdp committed 1 year ago

0e48b3d8

Optionally pass a `make_topology` for C API plugins.

zacmustin committed 1 year ago

9ddd2a36

Provide a default `axis_name` for `jax.pmap`.

danielsuo committed 1 year ago

b7c513fd

[direct-linearize] add direct linearize rule for cond

mattjj committed 1 year ago

45a3fe62

Older