Code

Fix debug_nans regressions.#25519

Code

Open

emilyfertig wants to merge 7 commits intojax-ml:mainjax-ml/jax:mainfrom emilyfertig:debug-nansemilyfertig/jax:debug-nansCopy head branch name to clipboard

+210-75

Commits

Commits on Jan 16, 2025

• Fix debug_nans regressions.

  Show description for 4917931
  

  emilyfertig

  and

  mattjj

  committed
  4917931

  Copy full SHA for 4917931

• Fix grad(shard_map(...)) error message to print the line that NaNed.

  emilyfertig

  committed
  0eac7f2

  Copy full SHA for 0eac7f2

• Address reviewer comments.

  emilyfertig

  committed
  7780468

  Copy full SHA for 7780468

• Update grad(shard_map(...)) test.

  emilyfertig

  committed
  8cc4d57

  Copy full SHA for 8cc4d57

• Changes from Matt's diff.

  emilyfertig

  committed
  184f14e

  Copy full SHA for 184f14e

• Reinstate the call stack from eqn.source_info in the backward pass.

  emilyfertig

  committed
  1a9c301

  Copy full SHA for 1a9c301

• Temporary hack so we can merge first debug_nans PR.

  emilyfertig

  committed
  023ad2e

  Copy full SHA for 023ad2e