Repository: google/tunix
File: tunix/rl/common.py (Python)

Code before:
) -> jax.Array:
  """Computes the per-token log probabilities."""
  prompt_completion_ids, positions, attn_mask = process_ids(
      prompt_tokens, completion_tokens, pad_id, eos_id, completion_mask
  )
  per_token_logps = get_per_token_logps(
      model,
      input_tokens=prompt_completion_ids,
      positions=positions,
      attn_mask=attn_mask,
      logits_to_keep=completion_tokens.shape[1],
  )
  if stop_gradient:
    per_token_logps = jax.lax.stop_gradient(per_token_logps)
  return per_token_logps


@nnx.jit(static_argnames=("pad_id", "eos_id", "stop_gradient"))
def compute_per_token_logps_and_logits(
    model: nnx.Module,
    prompt_tokens: jax.Array,
    completion_tokens: jax.Array,
    pad_id: int,
    eos_id: int,
    stop_gradient: bool = True,
    completion_mask: jax.Array | None = None,
) -> tuple[jax.Array, jax.Array]:
  """Computes per-token log probabilities and returns the logits slice.

  Returns:
    per_token_logps: [B, T]
    logits_slice:    [B, T, V]
  """
  prompt_completion_ids, positions, attn_mask = process_ids(
      prompt_tokens, completion_tokens, pad_id, eos_id, completion_mask
  )
  per_token_logps, logits_slice = get_per_token_logps(
      model,
      input_tokens=prompt_completion_ids,
      positions=positions,
      attn_mask=attn_mask,
      logits_to_keep=completion_tokens.shape[1],
      return_logits=True,
  )
  if stop_gradient:
    per_token_logps = jax.lax.stop_gradient(per_token_logps)
    logits_slice = jax.lax.stop_gradient(logits_slice)
  return per_token_logps, logits_slice


@nnx.jit(static_argnames=("pad_id", "eos_id", "stop_gradient"))

Review comment (suggestion):
Let's move completion_mask to before `stop_gradient` in the arg order. Let's do this in other fns too

Code after:
      positions=positions,
      attn_mask=attn_mask,
      logits_to_keep=completion_tokens.shape[1],
  )
  if stop_gradient:
    per_token_logps = jax.lax.stop_gradient(per_token_logps)
    logits = jax.lax.stop_gradient(logits)
  return per_token_logps, logits


@nnx.jit(static_argnames=("pad_id", "eos_id", "stop_gradient"))
def compute_score(
    model,
    prompt_tokens: jax.Array,
    completion_tokens: jax.Array,
    pad_id: int,
    eos_id: int,
    completion_mask: jax.Array | None = None,
    stop_gradient: bool = True,
):
  """Computes reward using the provided model."""
  prompt_completion_ids, positions, attn_mask = process_ids(
      prompt_tokens, completion_tokens, pad_id, eos_id, completion_mask
  )

  out = model(
      prompt_completion_ids,
      positions=positions,
      cache=None,
      attention_mask=attn_mask,
  )
  per_token_scores = out[0] if isinstance(out, tuple) else out
  # The model returns a tensor of shape [B, T, 1]. We squeeze the last
  # dimension to get a tensor of shape [B, T].
  per_token_scores = jnp.squeeze(per_token_scores, axis=-1)

  if stop_gradient:
    per_token_scores = jax.lax.stop_gradient(per_token_scores)

  return per_token_scores


def make_completion_mask(completion_ids, eos_tok: int = 0):
  """Create completion mask based on the EOS token.

  Args:
    completion_ids: Completion ids.
    eos_tok: EOS token id.

  Returns:
    Completion mask.