Use raval_pytree to calculate gradient norm to make independent of array sizes

JanLuca · JanLuca · commit 7be380b79cc4 · 2025-11-21T10:00:51.000+01:00
diff --git a/varipeps/optimization/optimizer.py b/varipeps/optimization/optimizer.py
@@ -844,7 +844,7 @@ def random_noise(a):
                             contraction = "precondition_operator"
 
                         grad_norm_squared = 1e-2 * (
-                            jnp.linalg.norm(jnp.asarray(working_gradient)) ** 2
+                            jnp.linalg.norm(ravel_pytree(working_gradient)[0]) ** 2
                         )
 
                         tmp_descent_dir = [

Original file line number	Diff line number	Diff line change
`@@ -844,7 +844,7 @@ def random_noise(a):`
`844`	`844`	`contraction = "precondition_operator"`
`845`	`845`
`846`	`846`	`grad_norm_squared = 1e-2 * (`
`847`		`- jnp.linalg.norm(jnp.asarray(working_gradient)) ** 2`
	`847`	`+ jnp.linalg.norm(ravel_pytree(working_gradient)[0]) ** 2`
`848`	`848`	`)`
`849`	`849`
`850`	`850`	`tmp_descent_dir = [`