style: improve func name for error tracing

Signed-off-by: nstarman <[email protected]>

Author: nstarman

quax/examples/zero/_core.py

@@ -61,7 +61,9 @@ def _(value: Zero, *, broadcast_dimensions, shape, sharding=None) -> Zero:
 
 
 @quax.register(lax.convert_element_type_p)
-def _(value: Zero, *, new_dtype, weak_type, sharding=None) -> Zero:
+def convert_element_type_zero(
+    value: Zero, *, new_dtype, weak_type, sharding=None
+) -> Zero:
     # sharding was added around JAX 0.4.31, it seems.
     del weak_type, sharding
     return Zero(value.shape, new_dtype)

tests/myarray.py

@@ -368,11 +368,8 @@ def conv_general_dilated_p(
 
 
 @register(lax.convert_element_type_p)
-def convert_element_type_p(operand: MyArray, **kw: Any) -> MyArray:
-    return replace(
-        operand,
-        array=lax.convert_element_type_p.bind(operand.array, **kw),
-    )
+def convert_element_type_myarray(operand: MyArray, **kw: Any) -> MyArray:
+    return replace(operand, array=lax.convert_element_type_p.bind(operand.array, **kw))
 
 
 # ==============================================================================