Improve error message with class name.

Flax Team · Flax Authors · commit 712eb7ec2237 · 2025-09-09T02:21:05.000-07:00
PiperOrigin-RevId: 804790020
diff --git a/flax/nnx/bridge/wrappers.py b/flax/nnx/bridge/wrappers.py
@@ -226,7 +226,7 @@ def _get_module_method(module, method: tp.Callable[..., Any] | str | None):
   if not callable(method):
     class_name = type(module).__name__
     raise TypeError(
-      f"'{method}' must be a callable, got {type(method)}."
+      f"'{class_name}.{method}' must be a callable, got {type(method)}."
     )
 
   return method

Original file line number	Diff line number	Diff line change
`@@ -226,7 +226,7 @@ def _get_module_method(module, method: tp.Callable[..., Any] \| str \| None):`
`226`	`226`	`if not callable(method):`
`227`	`227`	`class_name = type(module).__name__`
`228`	`228`	`raise TypeError(`
`229`		`- f"'{method}' must be a callable, got {type(method)}."`
	`229`	`+ f"'{class_name}.{method}' must be a callable, got {type(method)}."`
`230`	`230`	`)`
`231`	`231`
`232`	`232`	`return method`