JAX 技能审计
- 作者仓库星标 39
- 作者仓库 awesome-omni-skill
JAX Skill
JAX is Autograd and XLA, brought together for high-performance machine learning research.
Contents
- Concepts & Theory
- Immutability
- The 4 Transformations
- Pytrees
- Code Examples
jit,grad,vmap,randomusage- Control Flow (
scan,cond,fori_loop) - Parallelism (
sharding)
Common Workflows
1. Developing a new Model
- Define your parameters as a Pytree (dict/dataclass).
- Define your forward pass function (pure).
- Define your loss function.
- Use
jax.value_and_gradto get gradients. - Use
jax.jitto speed up the update step. - See examples.md for snippets.
2. Debugging Shapes/NaNs
- Disable JIT:
jax.config.update("jax_disable_jit", True)to debug with standard python tools. - Use
jax.debug.printinside JITted functions.
- 流狐分类
- AI 智能
- 作者声明 Agent
- 未找到明确声明；不据此推断已兼容或已测试
- 静态检查
- 88 / 100 · 启发式扫描，不代表运行安全
- 作者 / 版本 / 许可
- @diegosouzapw · 未声明 license
- 流狐 Token 估算
- 低消耗
- 流狐接入估算
- 即装即用
- 是否需要外部 API Key
- 未发现要求
- 检测到的系统要求
- 未声明
- 底层运行要求
- Python
- 检测到的文件与系统行为
-
- 只读
- 允许写入 / 修改
- 检测到的网络行为
- 仅限本地
- 安装命令数
- 无（仅作为资料）
档案由构建时根据 SKILL.md 与安装命令自动衍生，可能与作者实际意图存在差异。
需要注意： 未限定 allowed-tools，默认拥有全部工具权限。
作者没有在当前 SKILL.md 中定义固定输出样例。 Contents
Concepts & Theory Immutability The 4 Transformations
Common Workflows
Common Workflows
1. Developing a new Model
Define your parameters as a Pytree (dict/dataclass). Define your forward pass function (pure). Define your loss function.
2. Debugging Shapes/NaNs
Disable JIT: jax.config.update("jaxdisablejit", True) to debug with standard python tools. Use jax.debug.print inside JITted functions.
# JAX Skill
JAX is Autograd and XLA, brought together for high-performance machine learning research.
## Contents
- [Concepts & Theory](reference.md)
- Immutability
- The 4 Transformations
- Pytrees
- [Code Examples](examples.md)
- `jit`, `grad`, `vmap`, `random` usage
- Control Flow (`scan`, `cond`, `fori_loop`)
- Parallelism (`sharding`)
## Common Workflows
### 1. Developing a new Model
1. Define your parameters as a Pytree (dict/dataclass).
2. Define your forward pass function (pure).
3. Define your loss function.
4. Use `jax.value_and_grad` to get gradients.
5. Use `jax.jit` to speed up the update step.
6. See [examples.md](examples.md) for snippets.
### 2. Debugging Shapes/NaNs
1. Disable JIT: `jax.config.update("jax_disable_jit", True)` to debug with standard python tools.
2. Use `jax.debug.print` inside JITted functions. 证据边界与执行链路
作者原文负责流程事实；流狐只索引当前章节、要点、文件与命令。
章节 -> Contents → Common Workflows → 1. Developing a new Model → 2. Debugging Shapes/NaNs
要点 -> JAX is Autograd and XLA, brought together for high-performance machine learning research.
文件/命令 -> jit · grad · vmap · random · scan · cond · foriloop · sharding
内容 SHA-256 -> 6dc651c2fe13
方法与流程
适用与边界
原文中的明确线索
jit、grad、vmap、random、scan、cond、foriloop、sharding