Merge pull request #275 from AI-Hypercomputer:log_prob_pytree_node

jetstream authors · jetstream authors · commit 29329e8e7382 · 2025-06-10T16:19:20.000-07:00
PiperOrigin-RevId: 769846778
diff --git a/jetstream/engine/engine_api.py b/jetstream/engine/engine_api.py
@@ -94,7 +94,6 @@ class ResultTokens(abc.ABC):
   )
   # log probabilities of the tokens. Shape: [batch, tokens]
   log_prob: Union[jax.Array, np.ndarray] = struct.field(
-      pytree_node=False,
       default=None,
   )
 

Original file line number	Diff line number	Diff line change
`@@ -94,7 +94,6 @@ class ResultTokens(abc.ABC):`
`94`	`94`	`)`
`95`	`95`	`# log probabilities of the tokens. Shape: [batch, tokens]`
`96`	`96`	`log_prob: Union[jax.Array, np.ndarray] = struct.field(`
`97`		`- pytree_node=False,`
`98`	`97`	`default=None,`
`99`	`98`	`)`
`100`	`99`