Adjust common block colors, fix README.

PiperOrigin-RevId: 660128213

README.md

@@ -153,11 +153,10 @@ intermediates = var.value
 ```
 
 To learn more about how to build and manipulate neural networks with Penzai,
-we recommend starting with the "How to Think in Penzai" tutorial ([V1 API version][how_to_think_1], [V2 API version][how_to_think_2]), or one
-of the other tutorials in the [Penzai documentation][].
+we recommend starting with the ["How to Think in Penzai" tutorial][how_to_think]
+or one of the other tutorials in the [Penzai documentation][].
 
-[how_to_think_1]: https://penzai.readthedocs.io/en/stable/notebooks/how_to_think_in_penzai.html
-[how_to_think_2]: https://penzai.readthedocs.io/en/stable/notebooks/how_to_think_in_penzai.html
+[how_to_think]: https://penzai.readthedocs.io/en/stable/notebooks/how_to_think_in_penzai.html
 [Penzai documentation]: https://penzai.readthedocs.io
 
 

penzai/models/transformer/model_parts.py

@@ -101,6 +101,10 @@ class TransformerFeedForward(pz.nn.Sequential):
 class TransformerBlock(pz.nn.Sequential):
   """Informatively-named Sequential subclass for the main transformer blocks."""
 
+  def treescope_color(self):
+    color = "oklch(0.785 0.103 186.9 / 1.0)"
+    return color, f"color-mix(in oklab, {color} 25%, white)"
+
 
 @pz.pytree_dataclass
 class TransformerLM(pz.nn.Layer):

penzai/nn/attention.py

@@ -245,6 +245,9 @@ def __call__(
     output = self.attn_value_to_output((attn, value), **side_inputs)
     return output
 
+  def treescope_color(self):
+    return "oklch(0.785 0.103 38.5 / 1.0)"
+
 
 @struct.pytree_dataclass
 class KVCachingAttention(layer_base.Layer):