{
  "source_file": "0930",
  "source_platform": "Hashnode",
  "author": "Wesley Kambale",
  "title_listed": "Building Neural Networks with Flax NNX",
  "keywords": "",
  "domain": "机器学习",
  "content_type": "深度技术解析",
  "is_paid_flag": "否",
  "url": "https://kambale.dev/flax-nnx",
  "publish_date": "2026-02-23",
  "corpus_stage": "",
  "notes": "",
  "id": "0a0dfde1fb97",
  "collected_at": "2026-10-02T14:11:05",
  "title": "Building Neural Networks with Flax NNX",
  "body_markdown": "Building Neural Networks with Flax NNX\nFlax NNX: Neural Networks That Feel Like PyTorch\nUpdated\nJuly 29, 2026\n•\n11\nmin read\n•\nView as Markdown\nW\nWesley Kambale\nI'm a Machine Learning Engineer passionate about building production-ready ML systems for the African market.\nWith experience in TensorFlow, Keras, and Python-based workflows, I help teams bridge the gap between machine learning research and real-world deployment—especially on resource-constrained devices.\nI'm also a Google Developer Expert in AI. I regularly speak at tech conferences including PyCon Africa, DevFest Kampala, DevFest Nairobi and more and also write technical articles on AI/ML here.\nPart of series\nJAX\nOver the past two weeks, we've learned that JAX is fast (\njit\n), that it eliminates loops (\nvmap\n), and that it computes gradients automatically (\ngrad\n). These are powerful primitives.\nBut if you've been following along, you might have noticed something uncomfortable: we've been passing arrays around manually, tracking parameters in tuples, and writing functions that return updated state. It works, but it doesn't feel like building neural networks. It feels like accounting.\nIf you've used PyTorch, you know how natural it is to define a model as a class, call\nmodel(x)\n, and let the framework handle the rest. That ergonomic experience is what made PyTorch the dominant framework for research.\nToday, we get that experience in JAX.\nFlax NNX\nis a neural network library that gives you PyTorch-style classes and methods while compiling down to JAX's XLA backend. You write object-oriented code. JAX runs functional code. NNX bridges the gap.\nBy the end of this article, we'll have:\nUnderstood why NNX exists and what problem it solves\nBuilt our first neural network using\nnnx.Module\nLearned how NNX handles the critical issue of random number generation\nCreated a CNN for image classification—the same architecture used in real production systems\nLet's build something real.\nThe State Problem\nTo understand why Flax NNX is necessary, we need to understand the fundamental tension between PyTorch and JAX.\nHow PyTorch Handles State\nIn PyTorch, a model is an object that contains its own parameters:\nimport torch.nn as nn\n\nclass PyTorchMLP(nn.Module):\n    def __init__(self):\n        super().__init__()\n        self.linear1 = nn.Linear(784, 128)\n        self.linear2 = nn.Linear(128, 10)\n    \n    def forward(self, x):\n        x = torch.relu(self.linear1(x))\n        return self.linear2(x)\n\nmodel = PyTorchMLP()\noutput = model(input_data)  # Parameters are hidden inside\nThis is intuitive. The\nnn.Linear\nlayers create their own weight matrices internally. You don't see them, you don't manage them, they just exist.\nBut this \"hidden state\" creates problems for JAX. When you call\njax.jit\non a function, JAX traces through it and compiles a computational graph. If that function secretly reads or writes to hidden variables, JAX can't see those operations, and it can't optimize them.\nHow Pure JAX Handles State\nJAX demands\npure functions\n: functions where the output depends only on the inputs, with no side effects. If you want to update parameters, you must pass them in and return them out:\ndef forward(params, x):\n    x = jax.nn.relu(x @ params['w1'] + params['b1'])\n    return x @ params['w2'] + params['b2']\n\ndef train_step(params, x, y):\n    loss, grads = jax.value_and_grad(loss_fn)(params, x, y)\n    new_params = update_params(params, grads)\n    return new_params, loss  # Must return the new state\n\n# You're always passing params around\nparams, loss = train_step(params, x_batch, y_batch)\nparams, loss = train_step(params, x_batch, y_batch)\nparams, loss = train_step(params, x_batch, y_batch)\nThis is explicit and JIT-compatible, but it's tedious. For a model with dozens of layers, managing nested dictionaries of parameters becomes error-prone.\nHow Flax NNX Solves This\nNNX gives you the PyTorch interface while secretly doing the JAX bookkeeping. You write classes with attributes. NNX intercepts those attributes, extracts the parameters when needed for JIT compilation, and updates them in place when you're done.\nfrom flax import nnx\n\nclass NNX_MLP(nnx.Module):\n    def __init__(self, rngs: nnx.Rngs):\n        self.linear1 = nnx.Linear(784, 128, rngs=rngs)\n        self.linear2 = nnx.Linear(128, 10, rngs=rngs)\n    \n    def __call__(self, x):\n        x = nnx.relu(self.linear1(x))\n        return self.linear2(x)\n\nmodel = NNX_MLP(rngs=nnx.Rngs(0))\noutput = model(input_data)  # Looks just like PyTorch\nThe syntax is nearly identical to PyTorch. But under the hood, NNX can extract the parameters, pass them through JAX transformations, and put them back—all without you writing a single line of state management code.\nBuilding Your First NNX Model\nLet's build a simple multi-layer perceptron step by step.\nThe Basic Structure\nEvery NNX model inherits from\nnnx.Module\n:\nfrom flax import nnx\nimport jax.numpy as jnp\n\nclass SimpleMLP(nnx.Module):\n    def __init__(self, hidden_dim: int, output_dim: int, *, rngs: nnx.Rngs):\n        # Define layers as attributes\n        self.linear1 = nnx.Linear(784, hidden_dim, rngs=rngs)\n        self.linear2 = nnx.Linear(hidden_dim, output_dim, rngs=rngs)\n    \n    def __call__(self, x):\n        # Define forward pass\n        x = self.linear1(x)\n        x = nnx.relu(x)\n        x = self.linear2(x)\n        return x\nKey differences from PyTorch:\nNo\nsuper().__init__()\n: NNX uses Python metaclasses, so you don't need to call the parent constructor.\n__call__\ninstead of\nforward\n: In Python,\n__call__\nmakes an object callable. NNX uses this standard convention rather than PyTorch's custom\nforward\nmethod.\nThe\nrngs\nparameter\n: This is required. We'll explain why in the next section.\nInstantiating the Model\n# Create a random number generator\nrngs = nnx.Rngs(42)  # Seed for reproducibility\n\n# Create the model\nmodel = SimpleMLP(hidden_dim=128, output_dim=10, rngs=rngs)\n\n# Test with dummy data\nx = jnp.ones((32, 784))  # Batch of 32, 784 features each\noutput = model(x)\n\nprint(f\"Input shape:  {x.shape}\")       # (32, 784)\nprint(f\"Output shape: {output.shape}\")  # (32, 10)\nThat's it. The model is ready to use.\nThe Randomness Requirement\nYou might be wondering: why do we need to pass\nrngs\neverywhere?\nThe Problem with Hidden Randomness\nIn NumPy or PyTorch, random number generation uses a global state:\nimport numpy as np\nnp.random.seed(42)\nprint(np.random.randn())  # Always the same value\nprint(np.random.randn())  # Different value (global state advanced)\nThis global state is invisible. It changes every time you call a random function. And as we established in Week 1, JAX cannot work with hidden, changing state.\nJAX's Explicit Randomness\nIn JAX, randomness is deterministic. You create a \"key,\" and that key always produces the same random numbers:\nfrom jax import random\n\nkey = random.PRNGKey(42)\nprint(random.normal(key, shape=(3,)))  # Always the same\nprint(random.normal(key, shape=(3,)))  # Still the same!\nTo get different random numbers, you must\nsplit\nthe key:\nkey, subkey = random.split(key)\nprint(random.normal(subkey, shape=(3,)))  # New values\nHow nnx.Rngs Helps\nWhen you create a neural network, every layer needs random numbers to initialize its weights. If you have 50 layers, that's 50 key splits you'd need to manage manually.\nnnx.Rngs\nautomates this. It's a key dispenser that splits and distributes keys automatically:\nrngs = nnx.Rngs(42)\n\n# When you create a layer, it asks rngs for keys internally\nlinear1 = nnx.Linear(10, 20, rngs=rngs)  # Gets its own keys\nlinear2 = nnx.Linear(20, 30, rngs=rngs)  # Gets different keys\n\n# Both layers have different, reproducible initializations\nThe critical benefit:\nreproducibility\n. If you and I both run\nnnx.Rngs(42)\n, we get identical models. This matters for debugging, for scientific reproducibility, and for distributed training where multiple machines must initialize the same model.\nInspecting Your Model\nPyTorch lets you\nprint(model)\nto see the architecture. NNX has something better:\nnnx.display()\n.\nnnx.display(model)\nThis produces a rich, hierarchical view showing:\nEvery layer and sublayer\nParameter shapes and dtypes\nThe total parameter count\nThe structure of the computational graph\nIn Jupyter notebooks, this renders as an interactive tree you can expand and collapse. It's invaluable for debugging shape mismatches and verifying your architecture.\nBuilding a CNN for Image Classification\nLet's build something more substantial: a convolutional neural network for classifying images. This is the same architecture pattern used in production image classifiers.\nThe Architecture\nWe'll build a classic CNN with:\nTwo convolutional blocks (Conv → ReLU → Pool)\nA flatten operation\nTwo dense layers for classification\nfrom flax import nnx\nimport jax.numpy as jnp\nfrom functools import partial\n\nclass CNN(nnx.Module):\n    \"\"\"A simple CNN for image classification.\"\"\"\n    \n    def __init__(self, num_classes: int, *, rngs: nnx.Rngs):\n        # Convolutional layers\n        self.conv1 = nnx.Conv(\n            in_features=1,      # Input channels (1 for grayscale)\n            out_features=32,    # Output channels\n            kernel_size=(3, 3), # 3x3 filters\n            rngs=rngs\n        )\n        self.conv2 = nnx.Conv(\n            in_features=32,\n            out_features=64,\n            kernel_size=(3, 3),\n            rngs=rngs\n        )\n        \n        # Dense layers\n        # After two 2x2 pooling operations on 28x28 input: 28 → 14 → 7\n        # So we have 64 channels × 7 × 7 = 3136 features\n        self.linear1 = nnx.Linear(3136, 256, rngs=rngs)\n        self.linear2 = nnx.Linear(256, num_classes, rngs=rngs)\n        \n        # Pooling as a reusable operation\n        self.pool = partial(nnx.avg_pool, window_shape=(2, 2), strides=(2, 2))\n    \n    def __call__(self, x):\n        # Block 1: Conv → ReLU → Pool\n        x = self.conv1(x)\n        x = nnx.relu(x)\n        x = self.pool(x)\n        \n        # Block 2: Conv → ReLU → Pool\n        x = self.conv2(x)\n        x = nnx.relu(x)\n        x = self.pool(x)\n        \n        # Flatten: (batch, height, width, channels) → (batch, features)\n        x = x.reshape(x.shape[0], -1)\n        \n        # Classification head\n        x = nnx.relu(self.linear1(x))\n        x = self.linear2(x)\n        \n        return x\nTesting the Model\n# Initialize\nmodel = CNN(num_classes=10, rngs=nnx.Rngs(0))\n\n# Create dummy input: batch of 4 grayscale 28×28 images\n# Shape: (batch, height, width, channels)\ndummy_input = jnp.ones((4, 28, 28, 1))\n\n# Forward pass\noutput = model(dummy_input)\n\nprint(f\"Input shape:  {dummy_input.shape}\")  # (4, 28, 28, 1)\nprint(f\"Output shape: {output.shape}\")       # (4, 10)\n\n# Inspect the model structure\nnnx.display(model)\nThe output shape is\n(4, 10)\n—four images, ten class scores each. This is exactly what we'd feed into a softmax for classification.\nHow NNX Compiles with JAX\nHere's the magic: even though we're writing object-oriented code, we can still use JAX's transformations.\nThe @nnx.jit Decorator\nNNX provides its own versions of JAX transforms that understand NNX objects:\n@nnx.jit\ndef forward(model, x):\n    return model(x)\n\n# This is JIT-compiled, just like @jax.jit\noutput = forward(model, dummy_input)\nWhen you call this function, NNX:\nExtracts\nthe model's parameters into a pure JAX pytree\nTraces\nthe computation with those parameters\nCompiles\nthe trace with XLA\nUpdates\nthe model object with any changed state\nYou write familiar OOP code. JAX gets the pure functions it needs. Everyone wins.\nPreview: The Split/Merge Pattern\nUnder the hood, NNX uses two key operations:\n# Split: separate structure from state\ngraphdef, state = nnx.split(model)\n\n# graphdef: the \"blueprint\" of the model (static)\n# state: the actual parameter values (dynamic, a pytree)\n\n# Merge: reconstruct the model from structure and state\nreconstructed_model = nnx.merge(graphdef, state)\nYou rarely need to call these directly—\n@nnx.jit\nhandles it automatically. But understanding this pattern helps when you need to do advanced things like:\nSaving and loading checkpoints (Week 9)\nDistributing models across devices (Week 10)\nCustom training loops with fine-grained control\nCommon Layers Reference\nHere are the NNX equivalents of layers you know from PyTorch:\nPyTorch\nFlax NNX\nNotes\nnn.Linear\nnnx.Linear\nSame signature\nnn.Conv2d\nnnx.Conv\nUses\nin_features\n/\nout_features\nnn.BatchNorm2d\nnnx.BatchNorm\nTracks running stats automatically\nnn.LayerNorm\nnnx.LayerNorm\nSame behavior\nnn.Dropout\nnnx.Dropout\nRequires\nrngs\n, respects\ndeterministic\nflag\nnn.Embedding\nnnx.Embed\nFor token embeddings\nnn.MultiheadAttention\nnnx.MultiHeadAttention\nTransformer attention\nActivation functions aren't layers in NNX, they're just functions:\nx = nnx.relu(x)\nx = nnx.gelu(x)\nx = nnx.softmax(x)\nx = nnx.sigmoid(x)\nWorking with Parameters\nSometimes you need direct access to the parameters—for logging, for custom initialization, or for freezing layers.\nAccessing Parameters\n# Access a specific layer's weights\nprint(model.linear1.kernel.value.shape)  # (3136, 256)\nprint(model.linear1.bias.value.shape)    # (256,)\n\n# Parameters are nnx.Param objects; .value gets the JAX array\nExtracting All Parameters\n# Get all parameters as a state object\nstate = nnx.state(model)\n\n# Or specifically just the trainable parameters\nparams = nnx.state(model, nnx.Param)\nUpdating Parameters\n# Update the model with new state\nnnx.update(model, new_state)\nThis becomes important next week when we build training loops.\nExercises\nBefore moving on, try these:\nAdd Dropout\n: Modify the CNN to include\nnnx.Dropout(rate=0.5, rngs=rngs)\nbetween the dense layers. Note that dropout needs its own RNG stream for the random mask.\nBuild a deeper network\n: Create a 4-layer MLP with hidden dimensions [512, 256, 128, 64]. Use a loop in\n__init__\nto avoid repetition.\nParameter counting\n: Write a function that takes an NNX model and returns the total number of trainable parameters. Hint: use\nnnx.state(model, nnx.Param)\nand\njax.tree_util.tree_map\n.\nQuick Reference\nfrom flax import nnx\nimport jax.numpy as jnp\n\n# Define a Model\nclass MyModel(nnx.Module):\n    def __init__(self, *, rngs: nnx.Rngs):\n        self.linear = nnx.Linear(10, 5, rngs=rngs)\n    \n    def __call__(self, x):\n        return nnx.relu(self.linear(x))\n\n# Instantiate\nmodel = MyModel(rngs=nnx.Rngs(0))\n\n# Forward Pass\noutput = model(jnp.ones((32, 10)))\n\n# Inspect\nnnx.display(model)\n\n# Access Parameters\nweights = model.linear.kernel.value\nstate = nnx.state(model, nnx.Param)\n\n# JIT Compile\n@nnx.jit\ndef forward(model, x):\n    return model(x)\n\n# Common Layers\nnnx.Linear(in_features, out_features, rngs=rngs)\nnnx.Conv(in_features, out_features, kernel_size, rngs=rngs)\nnnx.BatchNorm(num_features, rngs=rngs)\nnnx.Dropout(rate, rngs=rngs)\nnnx.Embed(num_embeddings, features, rngs=rngs)\nWhat's Next\nWe have a model. But a model that can't learn is just a random number generator with extra steps.\nNext week, we build the\ntraining loop\n. We'll use:\nnnx.value_and_grad\nto compute loss and gradients\nnnx.Optimizer\nto manage parameter updates\noptax\nto define the optimization algorithm\nMetrics to track accuracy and loss\nWe'll train our CNN on real data and watch the loss curve drop. That's when this all becomes real.\n#\njax\n#\nflax\n#\nneural-networks\n#\nnnx",
  "metadata": {
    "title": "Building Neural Networks with Flax NNX",
    "description": "Over the past two weeks, we've learned that JAX is fast (jit), that it eliminates loops (vmap), and that it computes gradients automatically (grad). These are powerful primitives.\nBut if you've been f",
    "keywords": null,
    "author": null,
    "og:title": "Building Neural Networks with Flax NNX",
    "og:description": "Over the past two weeks, we've learned that JAX is fast (jit), that it eliminates loops (vmap), and that it computes gradients automatically (grad). These are powerful primitives.\nBut if you've been f",
    "og:url": "https://kambale.dev/flax-nnx",
    "og:site_name": "Wesley Kambale",
    "og:image": "https://cloudmate-test.s3.us-east-1.amazonaws.com/uploads/covers/61143e31119030192497a888/476b79f9-b003-4338-86d7-9b270d60f550.png",
    "og:image:width": "1200",
    "og:image:height": "630",
    "og:type": "article",
    "twitter:card": "summary_large_image",
    "twitter:creator": "https://twitter.com/WesKambale",
    "twitter:title": "Building Neural Networks with Flax NNX",
    "twitter:description": "Over the past two weeks, we've learned that JAX is fast (jit), that it eliminates loops (vmap), and that it computes gradients automatically (grad). These are powerful primitives.\nBut if you've been f",
    "twitter:image": "https://cloudmate-test.s3.us-east-1.amazonaws.com/uploads/covers/61143e31119030192497a888/476b79f9-b003-4338-86d7-9b270d60f550.png",
    "article:published_time": "2026-02-23T19:46:48.189Z",
    "article:modified_time": "2026-07-29T00:59:23.047Z",
    "article:author": "https://kambale.dev/@kambale"
  },
  "paywall": {
    "listed_paid": false,
    "detected": false,
    "marker": "",
    "public_only": false,
    "policy": "付费/会员内容只保留页面公开可见部分，不绕过付费墙"
  },
  "crawl": {
    "success": true,
    "status_code": 200,
    "error": ""
  },
  "code_blocks": [
    {
      "file": "attachments/0a0dfde1fb97/code/body-000.txt",
      "chars": 371,
      "language": "text"
    },
    {
      "file": "attachments/0a0dfde1fb97/code/body-001.txt",
      "chars": 509,
      "language": "text"
    },
    {
      "file": "attachments/0a0dfde1fb97/code/body-002.txt",
      "chars": 388,
      "language": "text"
    },
    {
      "file": "attachments/0a0dfde1fb97/code/body-003.txt",
      "chars": 480,
      "language": "text"
    },
    {
      "file": "attachments/0a0dfde1fb97/code/body-004.txt",
      "chars": 367,
      "language": "text"
    },
    {
      "file": "attachments/0a0dfde1fb97/code/body-005.txt",
      "chars": 156,
      "language": "text"
    },
    {
      "file": "attachments/0a0dfde1fb97/code/body-006.txt",
      "chars": 163,
      "language": "text"
    },
    {
      "file": "attachments/0a0dfde1fb97/code/body-007.txt",
      "chars": 87,
      "language": "text"
    },
    {
      "file": "attachments/0a0dfde1fb97/code/body-008.txt",
      "chars": 265,
      "language": "text"
    },
    {
      "file": "attachments/0a0dfde1fb97/code/body-009.txt",
      "chars": 19,
      "language": "text"
    },
    {
      "file": "attachments/0a0dfde1fb97/code/body-010.txt",
      "chars": 1549,
      "language": "text"
    },
    {
      "file": "attachments/0a0dfde1fb97/code/body-011.txt",
      "chars": 409,
      "language": "text"
    },
    {
      "file": "attachments/0a0dfde1fb97/code/body-012.txt",
      "chars": 133,
      "language": "text"
    },
    {
      "file": "attachments/0a0dfde1fb97/code/body-013.txt",
      "chars": 288,
      "language": "text"
    },
    {
      "file": "attachments/0a0dfde1fb97/code/body-014.txt",
      "chars": 70,
      "language": "text"
    },
    {
      "file": "attachments/0a0dfde1fb97/code/body-015.txt",
      "chars": 204,
      "language": "text"
    },
    {
      "file": "attachments/0a0dfde1fb97/code/body-016.txt",
      "chars": 150,
      "language": "text"
    },
    {
      "file": "attachments/0a0dfde1fb97/code/body-017.txt",
      "chars": 63,
      "language": "text"
    },
    {
      "file": "attachments/0a0dfde1fb97/code/body-018.txt",
      "chars": 787,
      "language": "text"
    }
  ],
  "images": [
    {
      "status": "ok",
      "url": "https://cloudmate-test.s3.us-east-1.amazonaws.com/uploads/covers/61143e31119030192497a888/476b79f9-b003-4338-86d7-9b270d60f550.png",
      "bytes": 81420,
      "reused": true,
      "local_path": "attachments/0a0dfde1fb97/images/000.png",
      "src": "https://cloudmate-test.s3.us-east-1.amazonaws.com/uploads/covers/61143e31119030192497a888/476b79f9-b003-4338-86d7-9b270d60f550.png",
      "alt": "Building Neural Networks with Flax NNX",
      "downloaded": true
    }
  ],
  "audio": [],
  "videos": [],
  "links": [
    {
      "href": "https://kambale.dev/flax-nnx.md",
      "label": "View as Markdown",
      "type": "a"
    },
    {
      "href": "https://hashnode.com/@kambale",
      "label": "W",
      "type": "a"
    },
    {
      "href": "https://twitter.com/WesKambale",
      "label": "",
      "type": "a"
    },
    {
      "href": "https://linkedin.com/in/WesKambale",
      "label": "",
      "type": "a"
    },
    {
      "href": "https://kambale.dev/series/jax-series",
      "label": "Part of series JAX",
      "type": "a"
    },
    {
      "href": "https://kambale.dev/tag/jax",
      "label": "# jax",
      "type": "a"
    },
    {
      "href": "https://kambale.dev/tag/flax",
      "label": "# flax",
      "type": "a"
    },
    {
      "href": "https://kambale.dev/tag/neural-networks",
      "label": "# neural-networks",
      "type": "a"
    },
    {
      "href": "https://kambale.dev/tag/nnx",
      "label": "# nnx",
      "type": "a"
    }
  ],
  "repair_version": 2,
  "repaired_at": "2026-10-14T10:50:47Z",
  "source_url": "https://kambale.dev/flax-nnx",
  "papers": [],
  "actual_video_files": [],
  "body_chars": 15232,
  "expected_body_images": 1,
  "discovery_status": "ok",
  "delivery_status": "ok",
  "missing_reason": "",
  "raw_html": "attachments/0a0dfde1fb97/pages/source.html",
  "local_html": "attachments/0a0dfde1fb97/pages/article.html",
  "body_base_path": "attachments/0a0dfde1fb97/pages/",
  "body_markdown_path": "attachments/0a0dfde1fb97/pages/article.md"
}