Stabilize and move penzai.experimental.v2 to the top level.

This change moves the V2 implementation of Penzai's neural net system
out of experimental and into the main penzai namespace, along with
the corresponding toolshed modules and example models.

For compatibility, the penzai.experimental.v2 aliases will remain, and
will redirect to the stable locations (implemented using dynamic stub
modules).

PiperOrigin-RevId: 656536074

docs/_include/_glue_figures.ipynb

@@ -22,7 +22,7 @@
         "import myst_nb\n",
         "\n",
         "import penzai\n",
-        "from penzai.experimental.v2 import pz\n",
+        "from penzai import pz\n",
         "\n",
         "import jax\n",
         "import jax.numpy as jnp\n",
@@ -53,8 +53,8 @@
       },
       "outputs": [],
       "source": [
-        "from penzai.experimental.v2.models.transformer.variants import gemma\n",
-        "from penzai.experimental.v2.models.transformer.variants import llamalike_common"
+        "from penzai.models.transformer.variants import gemma\n",
+        "from penzai.models.transformer.variants import llamalike_common"
       ]
     },
     {

docs/api/penzai.experimental.v2.rst

@@ -19,7 +19,7 @@ You can read about the V2 design in the guide :doc:`"How to Think in Penzai (v2
 
 To use the V2 API, we suggest importing the `pz` alias namespace from `penzai.experimental.v2.pz`: ::
 
-  from penzai.experimental.v2 import pz
+  from penzai import pz
 
 The rest of this page lists the main components used in the V2 API.
 

docs/guides/howto_reference.md

@@ -5,7 +5,7 @@ This notebook is a guide to accomplishing a variety of tasks with Penzai, using
 
 For this guide, we assume you have imported the experimental V2 version of the `pz` alias namespace:
 ```
-from penzai.experimental.v2 import pz
+from penzai import pz
 ```
 
 ## Visualization
@@ -223,7 +223,7 @@ Penzai's Gemma implementation includes a conversion utility that converts the ["
 ```python
 import kagglehub
 import orbax.checkpoint
-from penzai.experimental.v2.models.transformer import variants
+from penzai.models.transformer import variants
 
 weights_dir = kagglehub.model_download('google/gemma/Flax/7b')
 ckpt_path = os.path.join(weights_dir, '7b')
@@ -239,7 +239,7 @@ Penzai also includes re-implementations of the architectures used by [Llama](htt
 
 ```python
 import transformers
-from penzai.experimental.v2.models.transformer import variants
+from penzai.models.transformer import variants
 
 # To load a Llama model:
 hf_model = transformers.LlamaForCausalLM.from_pretrained(...)

docs/guides/v2_differences.md

@@ -202,8 +202,8 @@ from penzai.deprecated.v1.example_models import simple_mlp
 import penzai.toolshed
 
 # New
-from penzai.experimental.v2 import pz
-from penzai.experimental.v2.models import simple_mlp
+from penzai import pz
+from penzai.models import simple_mlp
 import penzai.experimental.v2.toolshed
 ```
 
@@ -381,7 +381,7 @@ model = gemma.model_core.GemmaTransformer.from_pretrained(flax_params_dict)
 # (model is an instance of GemmaTransformer)
 
 # New
-from penzai.experimental.v2.models.transformer import variants
+from penzai.models.transformer import variants
 model = variants.gemma.gemma_from_pretrained_checkpoint(flax_params_dict)
 # (model is an instance of TransformerLM)
 ```

notebooks/v2_how_to_think_in_penzai.ipynb

@@ -91,8 +91,8 @@
    "outputs": [],
    "source": [
     "import penzai\n",
-    "from penzai.experimental.v2 import pz\n",
-    "from penzai.experimental.v2.models import simple_mlp"
+    "from penzai import pz\n",
+    "from penzai.models import simple_mlp"
    ]
   },
   {
@@ -1277,7 +1277,7 @@
    },
    "outputs": [],
    "source": [
-    "from penzai.experimental.v2.toolshed import basic_training\n",
+    "from penzai.toolshed import basic_training\n",
     "import optax"
    ]
   },

notebooks/v2_induction_heads.ipynb

@@ -153,9 +153,9 @@
    "outputs": [],
    "source": [
     "import penzai\n",
-    "from penzai.experimental.v2 import pz\n",
+    "from penzai import pz\n",
     "\n",
-    "from penzai.experimental.v2.models import transformer"
+    "from penzai.models import transformer"
    ]
   },
   {

notebooks/v2_induction_heads_2B.ipynb

@@ -165,9 +165,9 @@
    "outputs": [],
    "source": [
     "import penzai\n",
-    "from penzai.experimental.v2 import pz\n",
+    "from penzai import pz\n",
     "\n",
-    "from penzai.experimental.v2.models import transformer"
+    "from penzai.models import transformer"
    ]
   },
   {

notebooks/v2_jitting_and_sharding.ipynb

@@ -135,7 +135,7 @@
    "outputs": [],
    "source": [
     "import penzai\n",
-    "from penzai.experimental.v2 import pz"
+    "from penzai import pz"
    ]
   },
   {
@@ -146,9 +146,9 @@
    },
    "outputs": [],
    "source": [
-    "from penzai.experimental.v2.models import transformer\n",
-    "from penzai.experimental.v2.models import simple_mlp\n",
-    "from penzai.experimental.v2.toolshed import basic_training"
+    "from penzai.models import transformer\n",
+    "from penzai.models import simple_mlp\n",
+    "from penzai.toolshed import basic_training"
    ]
   },
   {
@@ -610,7 +610,7 @@
    },
    "outputs": [],
    "source": [
-    "from penzai.experimental.v2.toolshed import jit_wrapper"
+    "from penzai.toolshed import jit_wrapper"
    ]
   },
   {
@@ -1181,7 +1181,7 @@
    },
    "outputs": [],
    "source": [
-    "from penzai.experimental.v2.toolshed import sharding_util"
+    "from penzai.toolshed import sharding_util"
    ]
   },
   {
@@ -1336,8 +1336,8 @@
    },
    "outputs": [],
    "source": [
-    "from penzai.experimental.v2.toolshed import sharding_util\n",
-    "from penzai.experimental.v2.models.transformer.variants import llamalike_common"
+    "from penzai.toolshed import sharding_util\n",
+    "from penzai.models.transformer.variants import llamalike_common"
    ]
   },
   {

notebooks/v2_lora_from_scratch.ipynb

@@ -142,7 +142,7 @@
    "outputs": [],
    "source": [
     "import penzai\n",
-    "from penzai.experimental.v2 import pz"
+    "from penzai import pz"
    ]
   },
   {
@@ -164,11 +164,11 @@
    },
    "outputs": [],
    "source": [
-    "from penzai.experimental.v2.models import transformer\n",
-    "from penzai.experimental.v2.models import simple_mlp\n",
+    "from penzai.models import transformer\n",
+    "from penzai.models import simple_mlp\n",
     "from penzai.toolshed import token_visualization\n",
-    "from penzai.experimental.v2.toolshed import basic_training\n",
-    "from penzai.experimental.v2.toolshed import jit_wrapper"
+    "from penzai.toolshed import basic_training\n",
+    "from penzai.toolshed import jit_wrapper"
    ]
   },
   {

penzai/experimental/v2/core/random_stream.py → penzai/core/random_stream.py

@@ -18,7 +18,7 @@
 
 import jax
 from penzai.core import struct
-from penzai.experimental.v2.core import variables
+from penzai.core import variables
 
 
 @struct.pytree_dataclass

penzai/experimental/v2/core/variables.py → penzai/core/variables.py

@@ -77,9 +77,9 @@
 
 import jax
 import jax.numpy as jnp
+from penzai.core import auto_order_types
 from penzai.core import selectors
 from penzai.core import struct
-from penzai.experimental.v2.core import auto_order_types
 
 
 T = TypeVar("T")

penzai/deprecated/v1/pz/__init__.py

@@ -66,4 +66,4 @@
 
 from . import de
 from . import nn
-from penzai.experimental.v2.pz import ts  # pylint: disable=g-bad-import-order
+from penzai.pz import ts  # pylint: disable=g-bad-import-order

penzai/deprecated/v1/toolshed/sharding_util.py

@@ -21,7 +21,7 @@
 
 import jax
 from penzai.deprecated.v1 import pz
-from penzai.experimental.v2.toolshed import sharding_util as sharding_util_v2
+from penzai.toolshed import sharding_util as sharding_util_v2
 
 PyTreeOfArrays = Any
 PyTreeOfNamedArrays = Any