Remove unnecessary allclose parameters for consistency.

google · Feb 24, 2025 · 0785521 · 0785521
1 parent b0d4555
commit 0785521
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/flax/nnx/nn/normalization.py b/flax/nnx/nn/normalization.py
@@ -891,7 +891,7 @@ class WeightNorm(nnx.Module):
 
     >>> w = model.normed_linear.layer_instance.kernel.value
     >>> col_norms = np.linalg.norm(np.array(w), axis=0)
-    >>> np.testing.assert_allclose(col_norms, np.ones(4), rtol=1e-5, atol=1e-5)
+    >>> np.testing.assert_allclose(col_norms, np.ones(4))
 
   Args:
     layer_instance: The layer instance to wrap.