Creating a NNX model with no bias results in a None being passed to t…

…he overview. PiperOrigin-RevId: 720549126
google · Jan 28, 2025 · 77b0602 · 77b0602
1 parent b4e47eb
commit 77b0602
Showing 1 changed file with 7 additions and 0 deletions.
diff --git a/clu/parameter_overview.py b/clu/parameter_overview.py
@@ -98,6 +98,13 @@ def count_parameters(params: _ParamsContainer) -> int:
 
 
 def _make_row(name, value) -> _ParamRow:
+  if value is None:
+    return _ParamRow(
+        name=name,
+        shape=(),
+        dtype="",
+        size=0,
+    )
   return _ParamRow(
       name=name,
       shape=value.shape,