Merge pull request #4349 from 8bitmp3:define-model-surgery

PiperOrigin-RevId: 691593549

Author: Flax Authors

docs_nnx/guides/surgery.ipynb

@@ -6,7 +6,7 @@
    "source": [
     "# Model surgery\n",
     "\n",
-    "In this guide, you will learn how to perform model surgery in Flax NNX using several real-world scenarios:\n",
+    "Model surgery is an act of making modifications on an existing neural network's building blocks and parameters, such as layer replacement, parameter or state manipulation, or even \"monkey patching\". In this guide, you will learn how to perform model surgery in Flax NNX using several real-world scenarios:\n",
     "\n",
     "* __Pythonic `nnx.Module` manipulation__: Using Pythonic ways to manipulate sub-`Module`s given a model.\n",
     "\n",

docs_nnx/guides/surgery.md

@@ -10,7 +10,7 @@ jupytext:
 
 # Model surgery
 
-In this guide, you will learn how to perform model surgery in Flax NNX using several real-world scenarios:
+Model surgery is an act of making modifications on an existing neural network's building blocks and parameters, such as layer replacement, parameter or state manipulation, or even "monkey patching". In this guide, you will learn how to perform model surgery in Flax NNX using several real-world scenarios:
 
 * __Pythonic `nnx.Module` manipulation__: Using Pythonic ways to manipulate sub-`Module`s given a model.
 

docs_nnx/nnx_basics.ipynb

@@ -247,7 +247,7 @@
    "source": [
     "### Model surgery\n",
     "\n",
-    "Flax [`nnx.Module`](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/module.html)s are mutable by default. This means that their structure can be changed at any time, which makes model surgery quite easy as any sub-`Module` attribute can be replaced with anything else, such as new `Module`s, existing shared `Module`s, `Module`s of different types, and so on. Moreover, [`nnx.Variable`](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/variables.html#flax.nnx.Variable)s can also be modified or replaced/shared.\n",
+    "Flax [`nnx.Module`](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/module.html)s are mutable by default. This means that their structure can be changed at any time, which makes [model surgery](https://flax.readthedocs.io/en/latest/guides/surgery.html) quite easy, as any sub-`Module` attribute can be replaced with anything else, such as new `Module`s, existing shared `Module`s, `Module`s of different types, and so on. Moreover, [`nnx.Variable`](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/variables.html#flax.nnx.Variable)s can also be modified or replaced/shared.\n",
     "\n",
     "The following example shows how to replace the `Linear` layers in the `MLP` model from the previous example with `LoraLinear` layers:"
    ]

docs_nnx/nnx_basics.md

@@ -143,7 +143,7 @@ In Flax, [`nnx.Dropout`](https://flax.readthedocs.io/en/latest/api_reference/fla
 
 ### Model surgery
 
-Flax [`nnx.Module`](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/module.html)s are mutable by default. This means that their structure can be changed at any time, which makes model surgery quite easy as any sub-`Module` attribute can be replaced with anything else, such as new `Module`s, existing shared `Module`s, `Module`s of different types, and so on. Moreover, [`nnx.Variable`](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/variables.html#flax.nnx.Variable)s can also be modified or replaced/shared.
+Flax [`nnx.Module`](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/module.html)s are mutable by default. This means that their structure can be changed at any time, which makes [model surgery](https://flax.readthedocs.io/en/latest/guides/surgery.html) quite easy, as any sub-`Module` attribute can be replaced with anything else, such as new `Module`s, existing shared `Module`s, `Module`s of different types, and so on. Moreover, [`nnx.Variable`](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/variables.html#flax.nnx.Variable)s can also be modified or replaced/shared.
 
 The following example shows how to replace the `Linear` layers in the `MLP` model from the previous example with `LoraLinear` layers:
 

docs_nnx/why.rst

@@ -279,7 +279,7 @@ is a simplified version of a ``BatchNorm`` layer that updates the mean and varia
 Model surgery
 ^^^^^^^^^^^^^
 
-In Flax Linen, model surgery has historically been challenging because of two reasons:
+In Flax Linen, `model surgery <https://flax.readthedocs.io/en/latest/guides/surgery.html>`_ has historically been challenging because of two reasons:
 
 1. Due to lazy initialization, it is not guaranteed that you can replace a sub-``Module`` with a new one.
 2. The parameter structure is separated from the ``flax.linen.Module`` structure, which means you have to manually keep them in sync.