logs31:Get logits and log_prob

steps

• revisit old code
  • if we can get logits easily
  • if we have log_prob already
• double check values above are reasonable.