Code

[ROCm] Implement RNN support#171

Code

Open

Ruturaj4 wants to merge 59 commits intorocm-mainROCm/jax:rocm-mainfrom ci_rnn_finalROCm/jax:ci_rnn_finalCopy head branch name to clipboard

+1,565-600

Commits

Commits on Dec 19, 2024

• support head size of 256

  Show description for 307ea87
  kaixih

  committed
  307ea87

Commits on Dec 23, 2024

• Fix a typo in documentation for pinv function.

  liblaf

  committed
  75b5654

• [pallas:mosaic_gpu] Updated the lowering following the changes in in Mosaic GPU internals

  Show description for 3e7f481
  

  superbobry

  authored and

  Google-ML-Automation

  committed
  3e7f481

• [pallas:mosaic_gpu] Reduced duplication between _ensure_fa and _ensure_ir_value

  Show description for a51d627
  

  superbobry

  authored and

  Google-ML-Automation

  committed
  a51d627

• [pallas:triton] Add support for lowering int4 load.

  Show description for 83e60a9
  

  chr1sj0nes

  authored and

  Google-ML-Automation

  committed
  83e60a9

• [mosaic_gpu] Include Mosaic GPU dialect fiels into jaxlib

  superbobry

  committed
  8987867

• Merge pull request jax-ml#25667 from superbobry:main

  Show description for 76a51f8
  Google-ML-Automation

  committed
  76a51f8

• Remove casting from jax.nn.one_hot

  Show description for cb10710
  jakevdp

  committed
  cb10710

• Use the right include for gmock and gtest

  Show description for 68ec202
  

  superbobry

  authored and

  Google-ML-Automation

  committed
  68ec202

• Merge pull request jax-ml#24607 from kaixih:support_head_size_256

  Show description for 704185e
  Google-ML-Automation

  committed
  704185e

• Merge pull request jax-ml#25662 from liblaf:main

  Show description for 6c85e54
  Google-ML-Automation

  committed
  6c85e54

• Internal: use a single registry for abstractify APIs

  jakevdp

  committed
  ccc3a29

• Merge pull request jax-ml#25651 from jakevdp:combine-abstractify

  Show description for 51b5102
  Google-ML-Automation

  committed
  51b5102

• Update XLA dependency to use revision

  Show description for 23965b7
  Google-ML-Automation

  committed
  23965b7

• changelog: link to api compatibility & python version docs

  jakevdp

  committed
  c206ae7

• Merge pull request jax-ml#25673 from jakevdp:changelog

  Show description for 1087623
  Google-ML-Automation

  committed
  1087623

• Bump actions/upload-artifact from 4.4.3 to 4.5.0

  Show description for 7da753e
  dependabot[bot]

  authored
  7da753e

• Finalize deprecation of some symbols from jax.lib.xla_client

  jakevdp

  committed
  40fe4b8

• Merge pull request jax-ml#25669 from jakevdp:undep

  Show description for c57b49c
  Google-ML-Automation

  committed
  c57b49c

• [Mosaic:TPU] Vreg-slice-aligned offset changes with scratch retiling

  Show description for 3c79b98
  

  tlongeri

  authored and

  Google-ML-Automation

  committed
  3c79b98

• Switch mlir bindings from pybind11 to nanobind

  Show description for b8091a4
  Google-ML-Automation

  committed
  b8091a4

• [Mosaic:TPU] In infer ext rule, avoid assigning offsets outside of dst first tile

  Show description for 4452960
  

  tlongeri

  authored and

  Google-ML-Automation

  committed
  4452960

Commits on Dec 24, 2024

• Merge pull request jax-ml#25674 from jax-ml:dependabot/github_actions/actions/upload-artifact-4.5.0

  Show description for fa9c7ed
  Google-ML-Automation

  committed
  fa9c7ed

• Merge pull request jax-ml#25672 from jakevdp:finalize-dep

  Show description for 4eff131
  Google-ML-Automation

  committed
  4eff131

• [pallas:mosaic_gpu] Addressed a todo in broadcasted_iota lowering

  Show description for 44333e1
  

  superbobry

  authored and

  Google-ML-Automation

  committed
  44333e1

• Update XLA dependency to use revision

  Show description for 64511a1
  Google-ML-Automation

  committed
  64511a1

Commits on Dec 25, 2024

• [AutoPGLE] Explicitly disable command buffers when profiler is used.

  Show description for b6aead6
  Google-ML-Automation

  committed
  b6aead6

• Update XLA dependency to use revision

  Show description for 42a0d55
  Google-ML-Automation

  committed
  42a0d55

Commits on Dec 26, 2024

• Fix formatting in the docs for transposing pytrees

  Mikcl

  committed
  008c25a

• Update XLA dependency to use revision

  Show description for 6dbda90
  Google-ML-Automation

  committed
  6dbda90

Commits on Dec 27, 2024

• Update XLA dependency to use revision

  Show description for 7ab61b7
  Google-ML-Automation

  committed
  7ab61b7

Commits on Dec 28, 2024

• [pallas:mosaic_gpu] Added some runtime type checking to copy_* and barrier_* primitives

  Show description for 76ccb19
  

  superbobry

  authored and

  Google-ML-Automation

  committed
  76ccb19

• Update XLA dependency to use revision

  Show description for 8eeedd1
  Google-ML-Automation

  committed
  8eeedd1

Commits on Dec 29, 2024

• Update XLA dependency to use revision

  Show description for 879fa12
  Google-ML-Automation

  committed
  879fa12

• Fixes the random key sharding in shard_map.

  yliu120

  committed
  97b1faa

Commits on Dec 30, 2024

• Update XLA dependency to use revision

  Show description for 25fff52
  Google-ML-Automation

  committed
  25fff52

• Merge pull request jax-ml#25692 from yliu120:rng_key_sharding

  Show description for 494c157
  Google-ML-Automation

  committed
  494c157

Commits on Dec 31, 2024

• Update XLA dependency to use revision

  Show description for e37ea58
  Google-ML-Automation

  committed
  e37ea58

Commits on Jan 1, 2025

• Fix log10 and log2 for large inputs.

  pearu

  committed
  50670bd

• Update XLA dependency to use revision

  Show description for 4a6cfeb
  Google-ML-Automation

  committed
  4a6cfeb

Commits on Jan 2, 2025

• tbp nightly instructions

  rdyro

  committed
  213e178

• Reverts 83e60a9

  Show description for dbe9ccd
  

  apaszke

  authored and

  Google-ML-Automation

  committed
  dbe9ccd

• Don't use x32 mode for pallas_test

  Show description for 7c984c6
  

  apaszke

  authored and

  Google-ML-Automation

  committed
  7c984c6

• Merge pull request jax-ml#25661 from rdyro:tb-nightly-instructions

  Show description for 04a0fbe
  Google-ML-Automation

  committed
  04a0fbe

• Merge pull request jax-ml#25706 from pearu:pearu/log10-large

  Show description for 82001ed
  Google-ML-Automation

  committed
  82001ed

• [Mosaic:TPU][NFC] Clean up unused variable

  Show description for ac817b4
  

  tlongeri

  authored and

  Google-ML-Automation

  committed
  ac817b4

• Fix OSS build for the Mosaic GPU dialect

  apaszke

  committed
  6443343

• Merge pull request jax-ml#25710 from apaszke:mgpu_dialect_fix

  Show description for 68483b8
  Google-ML-Automation

  committed
  68483b8

• Update XLA dependency to use revision

  Show description for 726950b
  Google-ML-Automation

  committed
  726950b

• Merge pull request jax-ml#25686 from Mikcl:docs/working-with-pytrees-formatting

  Show description for 800f903
  Google-ML-Automation

  committed
  800f903

• Compute cost-analysis on only one HLO module.

  Show description for df36c29
  

  zacmustin

  authored and

  Google-ML-Automation

  committed
  df36c29

• [Mosaic] NFC: Pull out vreg related functions to util.

  Show description for 57b2154
  

  WindQAQ

  authored and

  Google-ML-Automation

  committed
  57b2154

Commits on Jan 3, 2025

• jax.debug.print: respect local np.printoptions

  jakevdp

  committed
  3306063

• Merge pull request jax-ml#25713 from jakevdp:debug-printoptions

  Show description for 0f4677b
  Google-ML-Automation

  committed
  0f4677b

• Update XLA dependency to use revision

  Show description for e4278f7
  Google-ML-Automation

  committed
  e4278f7

Commits on Jan 4, 2025

• Update XLA dependency to use revision

  Show description for 9af2970
  Google-ML-Automation

  committed
  9af2970

Commits on Jan 5, 2025

• Add SMEM as a supported Pallas output memory space.

  Show description for 54fd738
  Google-ML-Automation

  committed
  54fd738

• Update XLA dependency to use revision

  Show description for d0a92c5
  Google-ML-Automation

  committed
  d0a92c5

• [ROCm] Implement RNN support

  Ruturaj4

  committed
  35939b9