nnx_functional_api__split_merge_state_update_graphdef_

nnx_functional_api__split_merge_state_update_graphdef_ is a cursor rule for Cursor from altaidevorg/rules-for-ai. It costs 33 tokens per session (4,767 once invoked), scanned A, original, MIT.

Functions that connect Flax NNX’s stateful Python modules with JAX’s function-based processing. They can separate a module’s structure from its current values, combine them again, inspect state, and apply updates.

In plain words
What is it for?
Use them when manually preparing NNX modules for JAX operations, inspecting or selecting state, rebuilding modules, or applying calculated updates.
Why use it?
JAX transformations work with explicit, unchanging data, while NNX modules hold changing values internally. These functions give you direct control over that conversion.

Cursor rule for Cursor

Written for Cursor: a Cursor rule (.mdc).

Install

Getting it into your agent

One page per mod, every tool's command on it. A separate URL per tool would split the same page into five that compete with each other.

agentmods
npx agentmods add rules/altaidevorg/rules-for-ai/nnx_functional_api__split_merge_state_update_graphdef_
Clone the repo
git clone --depth 1 https://github.com/altaidevorg/rules-for-ai

Made for: Cursor.

Wrote this? Show the measurements

A badge with what this costs and how it scanned, read live from this page, so it follows the numbers instead of freezing them. Markdown for a README, HTML for a documentation site or a project page.

agentmods badge for nnx_functional_api__split_merge_state_update_graphdef_

README.md
[![agentmods](https://agentmods.dev/badge/rules/altaidevorg/rules-for-ai/nnx_functional_api__split_merge_state_update_graphdef_.svg)](https://agentmods.dev/rules/altaidevorg/rules-for-ai/nnx_functional_api__split_merge_state_update_graphdef_)
Your own site
<a href="https://agentmods.dev/rules/altaidevorg/rules-for-ai/nnx_functional_api__split_merge_state_update_graphdef_"><img src="https://agentmods.dev/badge/rules/altaidevorg/rules-for-ai/nnx_functional_api__split_merge_state_update_graphdef_.svg" alt="Measured on agentmods" height="20"></a>
Per session 33 Only the description is in the session, so the agent can decide to use it. The body loads when it is invoked.
When invoked 4,767 The whole file, excluding the scripts and references it only reads on demand.
Security scan A 0 findings. Scan, not verified.
Origin original No closer match found in the catalogue.
Token cost

What it costs to keep this loaded

Counted locally with the o200k_base tokenizer, which is exact for GPT models; Claude uses its own tokenizer and its counts differ. Treat this as one consistent yardstick across the catalogue rather than a bill. Prices are per million input tokens.

ModelPer sessionOnce invoked
Fable 5.1 $0.00033 $0.04767
Opus 5 $0.00016 $0.02384
Sonnet 5 $0.00007 $0.00953
Haiku 4.5 $0.00003 $0.00477

Measured 5d ago against content hash 729b71f17034, method: parsed. Prices are Anthropic first-party input rates as of 2026-09-06, from the pricing page.

Security

Grade A, and why

nnx_functional_api__split_merge_state_update_graphdef_ scanned grade A with 0 findings against 26 rules in 11 categories — prompt injection, anti-refusal, data exfiltration, privilege escalation, supply chain, agent snooping, system-prompt leakage, SSRF and excessive agency — measured 5d ago.

A static scan of the body, not an audit. Every finding is printed with the line that produced it so you can judge whether it matters here. A mod is markdown that instructs an agent; that is exactly why what it instructs is worth reading.

Nothing flagged

None of the 26 patterns this scan looks for appear in this file: no shell pipes, no recursive deletes, no credential paths, no hidden text, no instruction-override or anti-refusal phrasing, no agent-config snooping. That is not a guarantee, it is the absence of the things that are checkable.

examples/flax/nnx_functional_api__split_merge_state_update_graphdef_.mdc · 334 lines

How it starts

The opening of the file, as written. The whole thing — 334 lines — stays where its author put it; the contents beside it link to each section on GitHub.

Chapter 6: NNX Functional API (split/merge/state/update/graphdef)

In the previous chapter, we saw how NNX Lifted Transforms like nnx.jit and nnx.grad automatically handle the state management required to apply JAX transformations to stateful nnx.Module methods. This chapter dives into the fundamental building blocks that make those lifted transforms possible: the NNX Functional API. This set of functions (nnx.split, nnx.merge, nnx.state, nnx.update, nnx.graphdef) provides the explicit bridge between NNX's object-oriented, stateful nnx.Module instances and JAX's functional programming model which requires pure functions operating on immutable pytrees.

Motivation: Manual Control over State for JAX Interoperability

JAX transformations operate on pure functions and immutable data structures (pytrees). nnx.Module instances, however, are stateful Python objects containing mutable nnx.Variable / nnx.VariableState attributes. To apply a JAX transform like jax.jit or jax.grad directly, we need to:

  1. Separate the static, structural information of the Module (its class types, attribute names, nested structure) from its dynamic state (the actual values held in Variables).
  2. Define a pure function that accepts the dynamic state as an explicit argument (a JAX pytree).
  3. Apply the JAX transformation to this pure function.
  4. Potentially recombine the static structure and the (possibly updated) dynamic state to get a Module instance back.

While Lifted Transforms automate this, the Functional API provides the manual controls (split, merge, etc.) to perform these steps explicitly. Understanding this API is crucial for advanced use cases, debugging, and comprehending how NNX integrates with JAX at a fundamental level.

Central Use Case: Manual JIT Compilation of a Module Method

Let's manually JIT-compile a method of an nnx.Module using the functional API. Compare this to the automated approach shown with nnx.jit in the previous chapter.

Read the full file on GitHub · 334 lines

Changes

What this file has done since we first saw it

Hashed on every crawl. A supply-chain change to an agent config is a question of when, not whether, so the history is kept rather than the latest state alone.

  1. 5d ago First seen · 334 lines · 33 tokens per session scan A 729b71f17034

Subscribe to this mod's changes

nnx_functional_api__split_merge_state_update_graphdef_ is a cursor rule published in the GitHub repository altaidevorg/rules-for-ai (2 stars, last pushed 1y ago), licensed MIT. It adds 33 tokens to every session and 4,767 once invoked, about $0.0002 per session on Opus 5. A static security scan graded it A with 0 findings. No closer match exists in the catalogue, so it is treated as the original; first seen 2026-08-31.