ROCm JAX Project Template

Warning

这个README是AI写的。

这是一个现代化的、高度自动化的 Python/JAX (ROCm AMD GPU) 项目模板。

通过精心组合当前最先进的开发工具链，本项目不仅使复杂的 AMD GPU 机器学习环境搭建变得前所未有的简单，还能够保证跨平台和团队成员之间完全一致的“开箱即用”体验。告别繁琐的环境配置环境，专注于核心代码开发。

🌟 核心工具链：现代与自动化的完美结合

本项目环境展开之所以既现代又高度自动化，得益于以下工具的协同工作：

• devenv.sh & Nix: 提供声明式、可重复的底层系统环境。无论是 Rust 编译器、Python 解释器，还是底层共享库（如 libffi, zlib, libdrm 等），Nix 都会精确、隔离地为你自动准备好，无需污染宿主机环境或手动编写复杂的安装脚本。
• uv: 采用 Rust 编写的极速 Python 包管理器。它无缝接管了 pyproject.toml，以惊人的速度完成依赖解析和虚拟环境的搭建。复杂如 jax[rocm7-local] 的依赖也能飞速且稳定地就绪。
• Devcontainer: 为 VS Code / Cursor 提供标准的容器化统一开发体验配置，结合底层 Nix 工具链，确保系统级配置和编辑器插件的零配置启动。
• JAX + ROCm: 强大的高性能数值计算框架，无缝运行在 AMD GPU 底层生态之上。

🚀 为什么选择本模板？

在传统的 ROCm 深度学习开发中，往往需要反复折腾环境变量、系统库版本 (LD_LIBRARY_PATH) 和 Python 依赖冲突。而在本模板中：

1. 环境自动展开: 开发容器启动或运行 devenv shell 的瞬间，Nix 和 uv 会自动在后台组装好包含底层 C/C++ 依赖和上层 Python 包的完整沙盒环境。
2. 极速冷启动: 摒弃了缓慢的传统包安装方式，利用 uv 的极速解析和并行下载，使环境部署时间缩短数十倍。
3. 强大的代码质量工具栏: 已经内置并自动配置好 treefmt (包含 nixfmt, rustfmt 等) 以及对应语言的 LSP 支持。
4. 完全可重现: 真正的“在我的机器上能跑，在你的机器上也能跑”，无需编写长篇幅的手动安装说明。

🛠️ 快速开始

方式一：使用 VS Code Devcontainer（推荐）

1. 在 VS Code 中安装 Dev Containers 插件。
2. 打开本项目所在的文件夹。
3. 按照右下角提示点击 Reopen in Container（或通过命令面板 Ctrl+Shift+P -> Dev Containers: Reopen in Container）。
4. 🌟 喝口水，环境构建自动完成，你可以直接运行 JAX 测试代码！

方式二：使用本地 devenv 环境

如果你偏好在纯系 Linux 环境中使用 Nix，比如你的 Linux 已经正确展开了ROCm：

1. 确保已安装了 Nix包管理器 和 devenv 工具。
2. 在项目根目录执行：

   devenv shell
   

   Nix 会自动获取 devenv.nix 中的依赖并进入开发环境。
3. 依赖就绪后，通过 uv 运行你的入口代码：

   uv run python src/rocm_jax/test_amdgpu.py