From 6917a15491544f68eade01195ba156fd91844409 Mon Sep 17 00:00:00 2001
From: Hana Joo <hanajoo@google.com>
Date: Tue, 7 Jul 2026 17:49:50 -0700
Subject: [PATCH] Automated Code Change

PiperOrigin-RevId: 944186032
---
 jmp/_src/loss_scale.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/jmp/_src/loss_scale.py b/jmp/_src/loss_scale.py
index 0bcdc31..b257315 100644
--- a/jmp/_src/loss_scale.py
+++ b/jmp/_src/loss_scale.py
@@ -116,11 +116,11 @@ class DynamicLossScale:
   ...   # conditionally update params using grads
   """
   loss_scale: jnp.ndarray
-  counter: jnp.ndarray = dataclasses.field(
+  counter: jnp.ndarray = dataclasses.field(  # pyrefly: ignore[bad-assignment]
       default_factory=lambda: np.zeros([], np.int32))
   period: int = 2000
   factor: int = 2
-  min_loss_scale: jnp.ndarray = dataclasses.field(
+  min_loss_scale: jnp.ndarray = dataclasses.field(  # pyrefly: ignore[bad-assignment]
       default_factory=lambda: np.ones([], np.float32))
 
   def __post_init__(self) -> None: