Skip to content

Commit

Permalink
Simplify default configuration for pz.de.TellIntermediate.
Browse files Browse the repository at this point in the history
PiperOrigin-RevId: 625558643
  • Loading branch information
danieldjohnson authored and Penzai Developers committed Apr 17, 2024
1 parent 40410b1 commit e26613d
Showing 1 changed file with 2 additions and 2 deletions.
4 changes: 2 additions & 2 deletions penzai/data_effects/side_output.py
Original file line number Diff line number Diff line change
Expand Up @@ -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))

0 comments on commit e26613d

Please sign in to comment.