Merge pull request #260 from AI-Hypercomputer:maxtext_logp

PiperOrigin-RevId: 751216668

Author: jetstream authors

CONTRIBUTING.md

@@ -25,6 +25,10 @@ This project follows
 
 ## Contribution process
 
+### Style Formatting
+
+Please run `make format` before submitting your PR. 
+
 ### Code Reviews
 
 All submissions, including submissions by project members, require review. We

jetstream/engine/engine_api.py

@@ -60,6 +60,7 @@ class SlotData:
   tokens: Union[jax.Array, np.ndarray]
   valid: Union[jax.Array, np.ndarray]
   lengths: Union[jax.Array, np.ndarray]
+  log_prob: Union[jax.Array, np.ndarray] = None
 
 
 # pylint: disable=g-doc-args
@@ -91,6 +92,11 @@ class ResultTokens(abc.ABC):
   samples_per_slot: int = struct.field(
       pytree_node=False,
   )
+  # log probabilities of the tokens. Shape: [batch, tokens]
+  log_prob: Union[jax.Array, np.ndarray] = struct.field(
+      pytree_node=False,
+      default=None,
+  )
 
   def copy_to_host_async(self: "ResultTokens") -> None:
     """Copy to host asynchronously."""
@@ -107,6 +113,7 @@ def convert_to_numpy(self: "ResultTokens") -> "ResultTokens":
         self.valid_idx,
         self.length_idx,
         self.samples_per_slot,
+        self.log_prob,
     )
 
   def get_result_at_slot(self, slot: int) -> SlotData:
@@ -148,6 +155,7 @@ def get_result_at_slots(self, slots: tuple[int]) -> SlotData:
         valid=self.data[slots, self.valid_idx[0] : self.valid_idx[1]],
         # Only get a 1D representation here
         lengths=self.data[slots, self.length_idx[0] : self.length_idx[1]][:, 0],
+        log_prob=self.log_prob[slots, :] if self.log_prob is not None else None,
     )