Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

[ROCm] Implement RNN support #171

Open
wants to merge 59 commits into
base: rocm-main
Choose a base branch
from
Open
Changes from 1 commit
Commits
Show all changes
59 commits
Select commit Hold shift + click to select a range
307ea87
support head size of 256
kaixih Oct 29, 2024
75b5654
Fix a typo in documentation for `pinv` function.
liblaf Dec 23, 2024
3e7f481
[pallas:mosaic_gpu] Updated the lowering following the changes in in …
superbobry Dec 23, 2024
a51d627
[pallas:mosaic_gpu] Reduced duplication between `_ensure_fa` and `_en…
superbobry Dec 23, 2024
83e60a9
[pallas:triton] Add support for lowering `int4` load.
chr1sj0nes Dec 23, 2024
8987867
[mosaic_gpu] Include Mosaic GPU dialect fiels into jaxlib
superbobry Dec 23, 2024
76a51f8
Merge pull request #25667 from superbobry:main
Google-ML-Automation Dec 23, 2024
cb10710
Remove casting from jax.nn.one_hot
jakevdp Dec 23, 2024
68ec202
Use the right include for gmock and gtest
superbobry Dec 23, 2024
704185e
Merge pull request #24607 from kaixih:support_head_size_256
Google-ML-Automation Dec 23, 2024
6c85e54
Merge pull request #25662 from liblaf:main
Google-ML-Automation Dec 23, 2024
ccc3a29
Internal: use a single registry for abstractify APIs
jakevdp Dec 23, 2024
51b5102
Merge pull request #25651 from jakevdp:combine-abstractify
Google-ML-Automation Dec 23, 2024
23965b7
Update XLA dependency to use revision
Google-ML-Automation Dec 23, 2024
c206ae7
changelog: link to api compatibility & python version docs
jakevdp Dec 23, 2024
1087623
Merge pull request #25673 from jakevdp:changelog
Google-ML-Automation Dec 23, 2024
7da753e
Bump actions/upload-artifact from 4.4.3 to 4.5.0
dependabot[bot] Dec 23, 2024
40fe4b8
Finalize deprecation of some symbols from jax.lib.xla_client
jakevdp Dec 23, 2024
c57b49c
Merge pull request #25669 from jakevdp:undep
Google-ML-Automation Dec 23, 2024
3c79b98
[Mosaic:TPU] Vreg-slice-aligned offset changes with scratch retiling
tlongeri Dec 23, 2024
b8091a4
Switch `mlir` bindings from `pybind11` to `nanobind`
Google-ML-Automation Dec 23, 2024
4452960
[Mosaic:TPU] In infer ext rule, avoid assigning offsets outside of ds…
tlongeri Dec 23, 2024
fa9c7ed
Merge pull request #25674 from jax-ml:dependabot/github_actions/actio…
Google-ML-Automation Dec 24, 2024
4eff131
Merge pull request #25672 from jakevdp:finalize-dep
Google-ML-Automation Dec 24, 2024
44333e1
[pallas:mosaic_gpu] Addressed a todo in `broadcasted_iota` lowering
superbobry Dec 24, 2024
64511a1
Update XLA dependency to use revision
Google-ML-Automation Dec 24, 2024
b6aead6
[AutoPGLE] Explicitly disable command buffers when profiler is used.
Google-ML-Automation Dec 25, 2024
42a0d55
Update XLA dependency to use revision
Google-ML-Automation Dec 25, 2024
008c25a
Fix formatting in the docs for transposing pytrees
Mikcl Dec 26, 2024
6dbda90
Update XLA dependency to use revision
Google-ML-Automation Dec 26, 2024
7ab61b7
Update XLA dependency to use revision
Google-ML-Automation Dec 27, 2024
76ccb19
[pallas:mosaic_gpu] Added some runtime type checking to `copy_*` and …
superbobry Dec 28, 2024
8eeedd1
Update XLA dependency to use revision
Google-ML-Automation Dec 28, 2024
879fa12
Update XLA dependency to use revision
Google-ML-Automation Dec 29, 2024
97b1faa
Fixes the random key sharding in shard_map.
yliu120 Dec 29, 2024
25fff52
Update XLA dependency to use revision
Google-ML-Automation Dec 30, 2024
494c157
Merge pull request #25692 from yliu120:rng_key_sharding
Google-ML-Automation Dec 30, 2024
e37ea58
Update XLA dependency to use revision
Google-ML-Automation Dec 31, 2024
50670bd
Fix log10 and log2 for large inputs.
pearu Jan 1, 2025
4a6cfeb
Update XLA dependency to use revision
Google-ML-Automation Jan 1, 2025
213e178
tbp nightly instructions
rdyro Dec 23, 2024
dbe9ccd
Reverts 83e60a9697ec20023f4e11169edf64e910b93031
apaszke Jan 2, 2025
7c984c6
Don't use x32 mode for pallas_test
apaszke Jan 2, 2025
04a0fbe
Merge pull request #25661 from rdyro:tb-nightly-instructions
Google-ML-Automation Jan 2, 2025
82001ed
Merge pull request #25706 from pearu:pearu/log10-large
Google-ML-Automation Jan 2, 2025
ac817b4
[Mosaic:TPU][NFC] Clean up unused variable
tlongeri Jan 2, 2025
6443343
Fix OSS build for the Mosaic GPU dialect
apaszke Jan 2, 2025
68483b8
Merge pull request #25710 from apaszke:mgpu_dialect_fix
Google-ML-Automation Jan 2, 2025
726950b
Update XLA dependency to use revision
Google-ML-Automation Jan 2, 2025
800f903
Merge pull request #25686 from Mikcl:docs/working-with-pytrees-format…
Google-ML-Automation Jan 2, 2025
df36c29
Compute cost-analysis on only one HLO module.
zacmustin Jan 2, 2025
57b2154
[Mosaic] NFC: Pull out vreg related functions to util.
WindQAQ Jan 2, 2025
3306063
jax.debug.print: respect local np.printoptions
jakevdp Jan 3, 2025
0f4677b
Merge pull request #25713 from jakevdp:debug-printoptions
Google-ML-Automation Jan 3, 2025
e4278f7
Update XLA dependency to use revision
Google-ML-Automation Jan 3, 2025
9af2970
Update XLA dependency to use revision
Google-ML-Automation Jan 4, 2025
54fd738
Add SMEM as a supported Pallas output memory space.
Google-ML-Automation Jan 5, 2025
d0a92c5
Update XLA dependency to use revision
Google-ML-Automation Jan 5, 2025
35939b9
[ROCm] Implement RNN support
Ruturaj4 Dec 2, 2024
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Prev Previous commit
Next Next commit
Use the right include for gmock and gtest
PiperOrigin-RevId: 709058082
superbobry authored and Google-ML-Automation committed Dec 23, 2024

Verified

This commit was created on GitHub.com and signed with GitHub’s verified signature. The key has expired.
commit 68ec202d452e24eaf0e9ab330d8fc3058ee85ec5
4 changes: 2 additions & 2 deletions jaxlib/mosaic/dialect/gpu/mosaic_gpu_test.cc
Original file line number Diff line number Diff line change
@@ -19,8 +19,8 @@ limitations under the License.
#include <optional>
#include <string>

#include "testing/base/public/gmock.h"
#include "testing/base/public/gunit.h"
#include <gmock/gmock.h>
#include <gtest/gtest.h>
#include "absl/container/flat_hash_set.h"
#include "absl/status/status.h"
#include "absl/status/statusor.h"