Merge pull request #3905 from chiamp:lm1b

PiperOrigin-RevId: 631512606
google · May 7, 2024 · 3d98eda · 3d98eda
2 parents 40c7eaa + fb5515c
commit 3d98eda
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 4 deletions.
diff --git a/examples/lm1b/train.py b/examples/lm1b/train.py
@@ -582,7 +582,7 @@ def encode_strings(strs, max_len):
           )  # pylint: disable=cell-var-from-loop
           summary["learning_rate"] = lr
           summary["perplexity"] = jnp.clip(
-              jnp.exp(summary["loss"]), a_max=1.0e4
+              jnp.exp(summary["loss"]), max=1.0e4
           )
           summary = {"train_" + k: v for k, v in summary.items()}
           writer.write_scalars(step, summary)
@@ -598,7 +598,7 @@ def encode_strings(strs, max_len):
           )
           # (clipped) perplexity after averaging log-perplexitie
           eval_results["perplexity"] = jnp.clip(
-              jnp.exp(eval_results["loss"]), a_max=1.0e4
+              jnp.exp(eval_results["loss"]), max=1.0e4
           )
           writer.write_scalars(
               step, {"eval_" + k: v for k, v in eval_results.items()}

diff --git a/flax/experimental/nnx/examples/lm1b/train.py b/flax/experimental/nnx/examples/lm1b/train.py
@@ -606,7 +606,7 @@ def constructor(config: models.TransformerConfig, key: jax.Array):
           )  # pylint: disable=cell-var-from-loop
           summary['learning_rate'] = lr
           summary['perplexity'] = jnp.clip(
-            jnp.exp(summary['loss']), a_max=1.0e4
+            jnp.exp(summary['loss']), max=1.0e4
           )
           summary = {'train_' + k: v for k, v in summary.items()}
           writer.write_scalars(step, summary)
@@ -621,7 +621,7 @@ def constructor(config: models.TransformerConfig, key: jax.Array):
           )
           # (clipped) perplexity after averaging log-perplexitie
           eval_results['perplexity'] = jnp.clip(
-            jnp.exp(eval_results['loss']), a_max=1.0e4
+            jnp.exp(eval_results['loss']), max=1.0e4
           )
           writer.write_scalars(
             step, {'eval_' + k: v for k, v in eval_results.items()}