Fix smm.py a_min parameter on jnp.clip()

QuinnMMcGarry · camiloCienet · commit ca749c8e4da7 · 2026-04-23T02:10:04.000Z
diff --git a/axlearn/common/ssm.py b/axlearn/common/ssm.py
@@ -122,7 +122,7 @@ def initialize(
                 * (math.log(cfg.dt_max) - math.log(cfg.dt_min))
                 + math.log(cfg.dt_min)
             )
-            dt = jnp.clip(dt, a_min=cfg.dt_init_floor)
+            dt = jnp.clip(dt, min=cfg.dt_init_floor)
             # Get inverse of softplus.
             inv_dt = dt + jnp.log(-jnp.expm1(-dt))
             return inv_dt

Original file line number	Diff line number	Diff line change
`@@ -122,7 +122,7 @@ def initialize(`
`122`	`122`	`* (math.log(cfg.dt_max) - math.log(cfg.dt_min))`
`123`	`123`	`+ math.log(cfg.dt_min)`
`124`	`124`	`)`
`125`		`- dt = jnp.clip(dt, a_min=cfg.dt_init_floor)`
	`125`	`+ dt = jnp.clip(dt, min=cfg.dt_init_floor)`
`126`	`126`	`# Get inverse of softplus.`
`127`	`127`	`inv_dt = dt + jnp.log(-jnp.expm1(-dt))`
`128`	`128`	`return inv_dt`