xla test

src/Compiler/Transform.idr

@@ -124,7 +124,7 @@ vmap n (MkFn params (Reshape from to y)) x =
 vmap n (MkFn params (Slice starts stops strides y)) x =
   Slice (0 :: starts) (n :: stops) (1 :: strides) (vmap n (MkFn params y) x)
 vmap n (MkFn params (DynamicSlice starts sizes y)) x =
-  -- DynamicSlice takes scalar arguments `starts`
+  -- takes scalar arguments
   let starts = (FromLiteral {dtype=U64} (Scalar Z) :: starts)
    in DynamicSlice starts (n :: sizes) (vmap n (MkFn params y) x)
 vmap n (MkFn params (Concat axis y z)) x =
@@ -138,8 +138,10 @@ vmap n (MkFn params (Broadcast {dtype} from to y)) x =
   Broadcast {dtype} (n :: from) (n :: to) (vmap n (MkFn params y) x)
 vmap n (MkFn params (Map f operands dimensions)) x = ?vmap_map
 vmap n (MkFn params (Reduce f neutral axes y)) x =
+  -- takes scalar arguments
   Reduce f neutral [| S axes |] (vmap n (MkFn params y) x)
 vmap n (MkFn params (Sort f dimension isStable ys)) x =
+  -- takes scalar arguments
   Sort f (S dimension) isStable (map (\op => vmap n (MkFn params op) x) ys)
 vmap n (MkFn params (Reverse axes y)) x = Reverse [| S axes |] (vmap n (MkFn params y) x)
 vmap n (MkFn params (Eq y z)) x = Eq (vmap n (MkFn params y) x) (vmap n (MkFn params z) x)

test/Unit/TestTensor.idr

@@ -1352,6 +1352,12 @@ normalIsReproducible = withTests 20 . property $ do
 
   sample ===# sample'
 
+xlaGraphs : Property
+xlaGraphs = fixedProperty $ do
+  let x = fromLiteral {dtype=S32} [0, 1, 2]
+      y = map (\x => fromLiteral (toLiteral x)) x
+  y ===# x
+
 export partial
 group : Group
 group = MkGroup "Tensor" $ [
@@ -1415,4 +1421,5 @@ group = MkGroup "Tensor" $ [
     , ("normal", normal)
     , ("normal updates seed", normalSeedIsUpdated)
     , ("normal produces same samples for same seed", normalIsReproducible)
+    , ("XLA Graph edge cases", xlaGraphs)
   ]