Fix return types: CBool for boolean ops, Complex for cplx/real/imag

- isZero, isInf, isNaN: Array a -> Array CBool (af_is* always emits u8)
- allTrue, anyTrue: Array a -> Int -> Array CBool (af_all/any_true emits u8)
- where': Array a -> Array Word32 (af_where emits u32 indices)
- cplx, cplx2, cplx2Batched: return Array (Complex a), not Array a
- real, imag: simplified to (RealFloat a, AFType a, AFType (Complex a))
  => Array (Complex a) -> Array a; previous signature was unlinked (a, b)
- Update tests to match corrected return types

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

Author: dmjio

src/ArrayFire/Algorithm.hs

@@ -27,6 +27,7 @@
 module ArrayFire.Algorithm where
 
 import Data.Word (Word32)
+import Foreign.C.Types (CBool)
 
 import ArrayFire.FFI
 import ArrayFire.Internal.Algorithm
@@ -154,13 +155,13 @@ max x (fromIntegral -> n) = x `op1` (\p a -> af_max p a n)
 -- [1 1 1 1]
 --         0
 allTrue
-  :: forall a. AFType a
+  :: AFType a
   => Array a
   -- ^ Array input
   -> Int
   -- ^ Dimension along which to see if all elements are True
-  -> Array a
-  -- ^ Will contain the maximum of all values in the input array along dim
+  -> Array CBool
+  -- ^ Will contain 1 where all elements along dim are true, 0 otherwise
 allTrue x (fromIntegral -> n) =
   x `op1` (\p a -> af_all_true p a n)
 
@@ -171,13 +172,13 @@ allTrue x (fromIntegral -> n) =
 -- [1 1 1 1]
 --         0
 anyTrue
-  :: forall a . AFType a
+  :: AFType a
   => Array a
   -- ^ Array input
   -> Int
-  -- ^ Dimension along which to see if all elements are True
-  -> Array a
-  -- ^ Returns if all elements are true
+  -- ^ Dimension along which to see if any elements are True
+  -> Array CBool
+  -- ^ Will contain 1 where any element along dim is true, 0 otherwise
 anyTrue x (fromIntegral -> n) =
   (x `op1` (\p a -> af_any_true p a n))
 
@@ -473,8 +474,8 @@ where'
   :: AFType a
   => Array a
   -- ^ Is the input array.
-  -> Array a
-  -- ^ will contain indices where input array is non-zero
+  -> Array Word32
+  -- ^ Indices where input array is non-zero
 where' = (`op1` af_where)
 
 -- | First order numerical difference along specified dimension.

src/ArrayFire/Arith.hs

@@ -28,7 +28,7 @@
 --------------------------------------------------------------------------------
 module ArrayFire.Arith where
 
-import Prelude                  (Bool(..), ($), (.), flip, fromEnum, fromIntegral, Real, RealFrac)
+import Prelude                  (Bool(..), ($), (.), flip, fromEnum, fromIntegral, Real, RealFloat)
 
 import Data.Coerce
 import Data.Proxy
@@ -1315,12 +1315,12 @@ atan2Batched x y (fromIntegral . fromEnum -> batch) = do
 --          (9.0000,9.0000)
 --          (10.0000,10.0000)
 cplx2
-  :: AFType a
+  :: (RealFloat a, AFType a, AFType (Complex a))
   => Array a
-  -- ^ First input
-  -> Array a
-  -- ^ Second input
+  -- ^ First input (real part)
   -> Array a
+  -- ^ Second input (imaginary part)
+  -> Array (Complex a)
   -- ^ Result of cplx2
 cplx2 x y =
   x `op2` y $ \arr arr1 arr2 ->
@@ -1342,14 +1342,14 @@ cplx2 x y =
 --          (9.0000,9.0000)
 --          (10.0000,10.0000)
 cplx2Batched
-  :: AFType a
+  :: (RealFloat a, AFType a, AFType (Complex a))
   => Array a
-  -- ^ First input
+  -- ^ First input (real part)
   -> Array a
-  -- ^ Second input
+  -- ^ Second input (imaginary part)
   -> Bool
   -- ^ Use batch
-  -> Array a
+  -> Array (Complex a)
   -- ^ Result of cplx2
 cplx2Batched x y (fromIntegral . fromEnum -> batch) = do
   x `op2` y $ \arr arr1 arr2 ->
@@ -1371,11 +1371,11 @@ cplx2Batched x y (fromIntegral . fromEnum -> batch) = do
 --          (9.0000,0.0000)
 --          (10.0000,0.0000)
 cplx
-  :: AFType a
+  :: (RealFloat a, AFType a, AFType (Complex a))
   => Array a
   -- ^ Input array
-  -> Array a
-  -- ^ Result of calling 'atan'
+  -> Array (Complex a)
+  -- ^ Complex array with input as real part and zero imaginary part
 cplx = flip op1 af_cplx
 
 -- | Execute real
@@ -1385,11 +1385,11 @@ cplx = flip op1 af_cplx
 -- [1 1 1 1]
 --    10.0000
 real
-  :: (AFType a, AFType (Complex b), RealFrac a, RealFrac b)
-  => Array (Complex b)
+  :: (RealFloat a, AFType a, AFType (Complex a))
+  => Array (Complex a)
   -- ^ Input array
   -> Array a
-  -- ^ Result of calling 'real'
+  -- ^ Real part of each element
 real = flip op1 af_real
 
 -- | Execute imag
@@ -1399,11 +1399,11 @@ real = flip op1 af_real
 -- [1 1 1 1]
 --    11.0000
 imag
-  :: (AFType a, AFType (Complex b), RealFrac a, RealFrac b)
-  => Array (Complex b)
+  :: (RealFloat a, AFType a, AFType (Complex a))
+  => Array (Complex a)
   -- ^ Input array
   -> Array a
-  -- ^ Result of calling 'imag'
+  -- ^ Imaginary part of each element
 imag = flip op1 af_imag
 
 -- | Execute conjg
@@ -2043,7 +2043,7 @@ isZero
   :: AFType a
   => Array a
   -- ^ Input array
-  -> Array a
+  -> Array CBool
   -- ^ Result of calling 'isZero'
 isZero = (`op1` af_iszero)
 
@@ -2066,7 +2066,7 @@ isInf
   :: (Real a, AFType a)
   => Array a
   -- ^ Input array
-  -> Array a
+  -> Array CBool
   -- ^ will contain 1's where input is Inf or -Inf, and 0 otherwise.
 isInf = (`op1` af_isinf)
 
@@ -2086,9 +2086,9 @@ isInf = (`op1` af_isinf)
 --          1
 --          1
 isNaN
-  :: forall a. (AFType a, Real a)
+  :: (AFType a, Real a)
   => Array a
   -- ^ Input array
-  -> Array a
+  -> Array CBool
   -- ^ Will contain 1's where input is NaN, and 0 otherwise.
 isNaN = (`op1` af_isnan)

test/ArrayFire/AlgorithmSpec.hs

@@ -94,13 +94,13 @@ spec =
       A.max (A.vector @(A.Complex Float) 3 [3 A.:+ 4, 1 A.:+ 0, 2 A.:+ 2]) 0 `shouldBe` A.scalar (3 A.:+ 4)
       A.max (A.vector @A.CBool 5 [0,1,1,0,1]) 0 `shouldBe` 1
     it "Should find if all elements are true along dimension" $ do
-      A.allTrue (A.vector @Double 5 (repeat 12.0)) 0 `shouldBe` 1
-      A.allTrue (A.vector @A.CBool 5 (repeat 1)) 0 `shouldBe` 1
-      A.allTrue (A.vector @A.CBool 5 (repeat 0)) 0 `shouldBe` 0
+      A.allTrue (A.vector @Double 5 (repeat 12.0)) 0 `shouldBe` A.scalar @A.CBool 1
+      A.allTrue (A.vector @A.CBool 5 (repeat 1)) 0 `shouldBe` A.scalar @A.CBool 1
+      A.allTrue (A.vector @A.CBool 5 (repeat 0)) 0 `shouldBe` A.scalar @A.CBool 0
     it "Should find if any elements are true along dimension" $ do
-      A.anyTrue (A.vector @A.CBool 5 (repeat 1)) 0 `shouldBe` 1
-      A.anyTrue (A.vector @Int 5 (repeat 23)) 0 `shouldBe` 1
-      A.anyTrue (A.vector @A.CBool 5 (repeat 0)) 0 `shouldBe` 0
+      A.anyTrue (A.vector @A.CBool 5 (repeat 1)) 0 `shouldBe` A.scalar @A.CBool 1
+      A.anyTrue (A.vector @Int 5 (repeat 23)) 0 `shouldBe` A.scalar @A.CBool 1
+      A.anyTrue (A.vector @A.CBool 5 (repeat 0)) 0 `shouldBe` A.scalar @A.CBool 0
     it "Should get count of all elements" $ do
       A.count (A.vector @Int 5 (repeat 1)) 0 `shouldBe` 5
       A.count (A.vector @A.CBool 5 (repeat 1)) 0 `shouldBe` 5
@@ -205,7 +205,7 @@ spec =
     describe "where'" $ do
       it "returns indices of nonzero elements" $ do
         A.where' (A.vector @Double 5 [0,1,0,2,0])
-          `shouldBe` A.vector @Double 2 [1,3]
+          `shouldBe` A.vector @A.Word32 2 [1,3]
       it "returns empty array when all elements are zero" $ do
         A.getDims (A.where' (A.vector @Double 3 [0,0,0]))
           `shouldBe` (0,1,1,1)

test/ArrayFire/ArithSpec.hs

@@ -140,15 +140,15 @@ spec =
       clamp (scalar @Int 2) (scalar @Int 1) (scalar @Int 3)
         `shouldBe` 2
     it "Should check if an array has positive or negative infinities" $ do
-      isInf (scalar @Double (1 / 0)) `shouldBe` scalar @Double 1
-      isInf (scalar @Double 10) `shouldBe` scalar @Double 0
+      isInf (scalar @Double (1 / 0)) `shouldBe` scalar @CBool 1
+      isInf (scalar @Double 10) `shouldBe` scalar @CBool 0
     it "Should check if an array has any NaN values" $ do
-      ArrayFire.isNaN (scalar @Double (acos 2)) `shouldBe` scalar @Double 1
-      ArrayFire.isNaN (scalar @Double 10) `shouldBe` scalar @Double 0
+      ArrayFire.isNaN (scalar @Double (acos 2)) `shouldBe` scalar @CBool 1
+      ArrayFire.isNaN (scalar @Double 10) `shouldBe` scalar @CBool 0
     it "Should check if an array has any Zero values" $ do
-      isZero (scalar @Double (acos 2)) `shouldBe` scalar @Double 0
-      isZero (scalar @Double 0) `shouldBe` scalar @Double 1
-      isZero (scalar @Double 1) `shouldBe` scalar @Double 0
+      isZero (scalar @Double (acos 2)) `shouldBe` scalar @CBool 0
+      isZero (scalar @Double 0) `shouldBe` scalar @CBool 1
+      isZero (scalar @Double 1) `shouldBe` scalar @CBool 0
 
     prop "Floating @Float (exp)" $ \(x :: Float) -> exp `shouldMatchBuiltin` exp $ x
     prop "Floating @Float (log)" $ \(x :: Float) -> log `shouldMatchBuiltin` log $ x