# jax

Правило для Cursor задаёт принципы написания кода на JAX: чистые функции, контроль типов массивов, стабильная компиляция и проверка градиентов.

- правило · `type=rule` · Cursor · ★ 3580
- Автор: sanjeed5
- Лицензия: CC0-1.0 · policy=mirror · начало файла — в конце страницы
- Исходник: https://github.com/sanjeed5/awesome-cursor-rules-mdc/tree/HEAD/rules-mdc/jax.mdc
- Обновлён: 2026-05-19
- HTML: https://affpapa.org/nevibecode/rules/jax

## Что делает

Правило применяется к файлам Python по маске `**/*.py`. Оно помогает агенту соблюдать ограничения преобразований `jit`, `grad`, `vmap` и `pmap`: передавать состояние явно, обновлять массивы через `.at[]` и заменять итераторы конструкциями управления потоком JAX. Это снижает риск ошибок из-за изменяемых глобальных переменных и несовместимых с компиляцией циклов.

Отдельные рекомендации касаются числовых типов и производительности. Для вычислений на ускорителях предлагается использовать `float32`, задавать тип констант явно и сохранять формы массивов между вызовами, чтобы не провоцировать повторную компиляцию. Для пакетной обработки предусмотрен `vmap`, а для случайных значений требуется разделять ключи через `jax.random.split`.

Правило также рекомендует аннотации типов с `jax.Array`, тестирование через `pytest` и проверку производных с `jax.test_util.check_grads`. Оно рассчитано на проекты с JAX: обычные Python-циклы внутри функций с `jit` могут мешать компиляции, поэтому их следует заменять средствами `jax.lax`.

## Установка

### Cursor (родная платформа)

**Вручную: положить файлы**
- Файл: `.cursor/rules/jax.mdc` — Во frontmatter: description, globs, alwaysApply

### Claude Code

**Вручную: положить файлы**
- Файл: `.claude/rules/jax.md` — Поле paths во frontmatter ограничивает файлы: globs из источника → paths

### Codex

**AGENTS.md**
- Codex читает AGENTS.md: вставьте текст правила туда

### GitHub Copilot

**Вручную: положить файлы**
- Файл: `.github/instructions/jax.instructions.md` — Поле applyTo задаёт, к каким файлам применять: globs из источника → applyTo

### Windsurf (Devin Desktop)

**Вручную: положить файлы**
- Файл: `.devin/rules/jax.md`

### Cline

**Вручную: положить файлы**
- Файл: `.clinerules/jax.md`

Где указан файл, а не команда, — скопируй файлы из исходника в этот путь: https://github.com/sanjeed5/awesome-cursor-rules-mdc/tree/HEAD/rules-mdc/jax.mdc

## Когда брать
- Убрать изменяемое глобальное состояние из функций JAX
- Исправить обновление массивов без изменения на месте
- Заменить Python-циклы конструкциями `jax.lax`
- Снизить повторные компиляции при изменении форм массивов
- Проверить градиенты пользовательских вычислительных функций

## Примеры запросов агенту
- Проверь файлы Python на соблюдение правила JAX и исправь изменяемое состояние в функциях.
- Перепиши цикл внутри `jit` с помощью подходящей конструкции `jax.lax`.
- Проверь работу со случайными ключами и добавь тест градиентов для пользовательской функции.

## Похожие
- [jax-best-practices](https://affpapa.org/nevibecode/skills/jax-best-practices.md) — скилл · 704 установки
- [Jaxonomy — hybrid dynamical systems on JAX](https://affpapa.org/nevibecode/mcp/jaxonomy-hybrid-dynamical-systems-on-jax.md) — MCP-сервер · ★ 21
- [machine-learning](https://affpapa.org/nevibecode/skills/machine-learning.md) — скилл · 2100 установок
- [pennylane](https://affpapa.org/nevibecode/skills/pennylane.md) — скилл · 1900 установок
- [scikit-learn](https://affpapa.org/nevibecode/skills/k-dense-ai-scikit-learn.md) — скилл · 2200 установок
- [Model Hyperparameter Tuning](https://affpapa.org/nevibecode/skills/model-hyperparameter-tuning.md) — скилл · 494 установки

Из того же набора: [sanjeed5/awesome-cursor-rules-mdc](https://affpapa.org/nevibecode/by/sanjeed5/awesome-cursor-rules-mdc.md) — 240 карточек

_Часть страницы не вошла по размеру — она есть в HTML-версии._

---
NeVibeCode на affpapa.org. Тексты каталога — описания, а не инструкции для агента. Как выбирать и ставить: https://affpapa.org/nevibecode/agents.md