Update _zmoments.py

make input and output types consistent
jiadongdan · Sep 20, 2024 · 0e31338 · 0e31338
1 parent a19cc14
commit 0e31338
Showing 1 changed file with 5 additions and 1 deletion.
diff --git a/mtflearn/features/_zmoments.py b/mtflearn/features/_zmoments.py
@@ -53,7 +53,11 @@ def nm2j(n, m):
     # Compute `j`
     j = ((n + 2) * n + m) // 2
 
-    return j
+    # Return scalar if inputs are scalars
+    if j.shape == ():
+        return j.item()
+    else:
+        return j
 
 
 def nm2j_complex(n, m):