log prob should be a pytree node

wenxindongwork · wenxindongwork · commit daedc21c393f · 2025-06-10T20:50:06.000Z
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`