From e78eecfb8d965ab48b8344dc599d4c7f8f9fa373 Mon Sep 17 00:00:00 2001
From: Daniel Johnson <ddjohnson@google.com>
Date: Thu, 24 Oct 2024 18:37:22 -0700
Subject: [PATCH] Remove remaining references to Penzai from NNX.

PiperOrigin-RevId: 689590143
---
 docs_nnx/guides/gemma.ipynb | 2 +-
 docs_nnx/guides/gemma.md    | 2 +-
 flax/nnx/graph.py           | 7 +++----
 3 files changed, 5 insertions(+), 6 deletions(-)

diff --git a/docs_nnx/guides/gemma.ipynb b/docs_nnx/guides/gemma.ipynb
index 7230ca1b02..1ac5dd12f4 100644
--- 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
index 69ff17acb5..1fa07ea3b2 100644
--- 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
index d7a8c1f47c..9b26495e70 100644
--- 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,