Remove remaining references to Penzai from NNX.

PiperOrigin-RevId: 689590143
google · Oct 25, 2024 · e78eecf · e78eecf
1 parent 40f080e
commit e78eecf
Show file tree

Hide file tree

Showing 3 changed files with 5 additions and 6 deletions.
diff --git a/docs_nnx/guides/gemma.ipynb b/docs_nnx/guides/gemma.ipynb
@@ -38,7 +38,7 @@
    "outputs": [],
    "source": [
     "! pip install --no-deps -U flax\n",
-    "! pip install jaxtyping kagglehub penzai"
+    "! pip install jaxtyping kagglehub treescope"
    ]
   },
   {

diff --git a/docs_nnx/guides/gemma.md b/docs_nnx/guides/gemma.md
@@ -30,7 +30,7 @@ You will find in this colab a detailed tutorial explaining how to use NNX to loa
 
 ```{code-cell} ipython3
 ! pip install --no-deps -U flax
-! pip install jaxtyping kagglehub penzai
+! pip install jaxtyping kagglehub treescope
 ```
 
 ## Downloading the checkpoint

diff --git a/flax/nnx/graph.py b/flax/nnx/graph.py
@@ -229,10 +229,9 @@ def __nnx_repr__(self):
     yield reprlib.Attr('type', self.type.__name__)
     yield reprlib.Attr('index', self.index)
 
-  def __penzai_repr__(self, path, subtree_renderer):
-    from penzai.treescope import repr_lib as pz_repr_lib  # type: ignore[import-not-found,import-untyped]
-
-    return pz_repr_lib.render_object_constructor(
+  def __treescope_repr__(self, path, subtree_renderer):
+    import treescope  # type: ignore[import-not-found,import-untyped]
+    return treescope.repr_lib.render_object_constructor(
       object_type=type(self),
       attributes={'type': self.type, 'index': self.index},
       path=path,