Simplify default configuration for pz.de.TellIntermediate.

PiperOrigin-RevId: 625558643
google-deepmind · Apr 17, 2024 · 2b4a40a · 2b4a40a
1 parent ed92dd0
commit 2b4a40a
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 3 deletions.
diff --git a/.github/workflows/unittests.yml b/.github/workflows/unittests.yml
@@ -18,7 +18,7 @@ jobs:
     # Install deps
     - uses: actions/setup-python@v4
       with:
-        python-version: 3.10
+        python-version: 3.10.14
         # Uncomment to cache of pip dependencies (if tests too slow)
         # cache: pip
         # cache-dependency-path: '**/pyproject.toml'

diff --git a/penzai/data_effects/side_output.py b/penzai/data_effects/side_output.py
@@ -223,13 +223,13 @@ def _make_ref(hole: SideOutputRequest):
 class TellIntermediate(layer_base.Layer):
   """Helper layer that writes its intermediate value to a side output."""
 
-  side_out: SideOutputEffect
+  side_out: SideOutputEffect = SideOutputRequest(tag="intermediate")
 
   def __call__(self, intermediate_value):
     self.side_out.tell(intermediate_value)
     return intermediate_value
 
   @classmethod
-  def from_config(cls, tag: SideOutputTag = "intermediate") -> TellIntermediate:
+  def from_config(cls, tag: SideOutputTag) -> TellIntermediate:
     """Builds a TellIntermediate layer that writes to the given tag."""
     return cls(side_out=SideOutputRequest(tag=tag))