From 9593c964185bf8132f6d0a2408d721c9b033aff8 Mon Sep 17 00:00:00 2001
From: Joel Berkeley <16429957+joelberkeley@users.noreply.github.com>
Date: Thu, 11 Jul 2024 03:06:47 +0100
Subject: [PATCH] wip

---
 spidr/src/Tensor.idr                        | 15 ++++------
 test/runner/TestRunner.idr                  |  2 +-
 test/runner/Unit/Model/TestKernel.idr       |  3 +-
 test/runner/Unit/TestDistribution.idr       |  5 +---
 test/runner/Unit/TestTensor.idr             | 25 ++---------------
 test/runner/Unit/TestTensor/Elementwise.idr | 31 +++++++--------------
 test/runner/Unit/TestTensor/HigherOrder.idr | 12 +-------
 test/runner/Unit/TestTensor/Sampling.idr    | 11 +-------
 test/runner/Unit/TestTensor/Slice.idr       |  9 +-----
 test/runner/Unit/TestTensor/Structure.idr   | 12 +-------
 test/runner/Utils.idr                       |  8 +++---
 test/runner/Utils/Comparison.idr            | 20 ++++++-------
 12 files changed, 39 insertions(+), 114 deletions(-)

diff --git a/spidr/src/Tensor.idr b/spidr/src/Tensor.idr
index 2ddcca840..15c1aa08f 100644
--- a/spidr/src/Tensor.idr
+++ b/spidr/src/Tensor.idr
@@ -132,9 +132,8 @@ namespace S32
   fromInteger : Integer -> Tensor [] S32
   fromInteger = tensor . Scalar . fromInteger
 
-partial
 try : Show e => EitherT e IO a -> IO a
-try = eitherT (idris_crash . show) pure
+try = eitherT (assert_total . idris_crash . show) pure
 
 namespace Tag
   ||| Evaluate a `Tensor`, returning its value as a `Literal`. This function builds and executes the
@@ -143,7 +142,7 @@ namespace Tag
   ||| **Note:** Each call to `eval` will rebuild and execute the graph; multiple calls to `eval` on
   ||| different tensors, even if they are in the same computation, will be treated independently.
   ||| To efficiently evaluate multiple tensors at once, use `TensorList.Tag.eval`.
-  export covering  -- is this true?
+  export covering
   eval : Device -> PrimitiveRW dtype ty => Tag (Tensor shape dtype) -> IO (Literal shape ty)
   eval device (MkTagT x) =
     let (env, MkTensor root) = runState empty x
@@ -152,12 +151,8 @@ namespace Tag
           [lit] <- execute device (MkFn [] root env) [shape]
           read {dtype} [] lit
 
--- is it safe to use this within a `Tag` context? It's probably not with `unsafePerformIO`, but
--- that's kind of expected. What about w/o `unsafePerformIO`?
 ||| A convenience wrapper for `Tag.eval`, for use with a bare `Tensor`.
-|||
-||| **Note:** It is not safe to use
-export covering  -- is this true?
+export covering
 eval : Device -> PrimitiveRW dtype ty => Tensor shape dtype -> IO (Literal shape ty)
 eval device x = eval device (pure x)
 
@@ -185,7 +180,7 @@ namespace TensorList
     ||| ```
     ||| In contrast to `Tensor.eval` when called on multiple tensors, this function constructs and
     ||| compiles the graph just once.
-    export covering  -- is this true?
+    export covering
     eval : Device -> Tag (TensorList shapes tys) -> IO (All2 Literal shapes tys)
     eval device (MkTagT xs) =
       let (env, xs) = runState empty xs
@@ -215,7 +210,7 @@ namespace TensorList
       readAll (MkTensor {dtype} _ :: ts) (l :: ls) = [| read {dtype} [] l :: readAll ts ls |]
 
   ||| A convenience wrapper for `TensorList.Tag.eval`, for use with a bare `TensorList`.
-  export covering  -- is this true?
+  export covering
   eval : Device -> TensorList shapes tys -> IO (All2 Literal shapes tys)
   eval device xs = eval device (pure xs)
 
diff --git a/test/runner/TestRunner.idr b/test/runner/TestRunner.idr
index e0a213150..792d2ae12 100644
--- a/test/runner/TestRunner.idr
+++ b/test/runner/TestRunner.idr
@@ -28,7 +28,7 @@ import Unit.TestTensor
 import Unit.TestLiteral
 import Unit.TestUtil
 
-export partial
+export
 run : Device -> IO ()
 run device = test [
       Utils.TestComparison.group
diff --git a/test/runner/Unit/Model/TestKernel.idr b/test/runner/Unit/Model/TestKernel.idr
index e75e36737..5e320759c 100644
--- a/test/runner/Unit/Model/TestKernel.idr
+++ b/test/runner/Unit/Model/TestKernel.idr
@@ -23,7 +23,6 @@ import Model.Kernel
 import Utils.Comparison
 import Utils.Cases
 
-partial
 rbfMatchesTFP : Device => Property
 rbfMatchesTFP = fixedProperty $ do
   let lengthScale = tensor 0.4
@@ -42,7 +41,7 @@ rbfMatchesTFP = fixedProperty $ do
         ]
   rbf lengthScale x x' ===# pure expected
 
-export partial
+export
 group : Device => Group
 group = MkGroup "Kernel" $ [
     ("rbf matches tfp", rbfMatchesTFP)
diff --git a/test/runner/Unit/TestDistribution.idr b/test/runner/Unit/TestDistribution.idr
index 8dc29e12e..67879ed49 100644
--- a/test/runner/Unit/TestDistribution.idr
+++ b/test/runner/Unit/TestDistribution.idr
@@ -23,7 +23,6 @@ import Distribution
 import Utils.Comparison
 import Utils.Cases
 
-partial
 gaussianUnivariatePDF : Device => Property
 gaussianUnivariatePDF = property $ do
   let doubles = literal [] doubles
@@ -37,7 +36,6 @@ gaussianUnivariatePDF = property $ do
     univariate : Double -> Double -> Double -> Double
     univariate x mean cov = exp (- (x - mean) * (x - mean) / (2 * cov)) / sqrt (2 * pi * cov)
 
-partial
 gaussianMultivariatePDF : Device => Property
 gaussianMultivariatePDF = fixedProperty $ do
   let mean = tensor [[-0.2], [0.3]]
@@ -45,7 +43,6 @@ gaussianMultivariatePDF = fixedProperty $ do
       x = tensor [[1.1], [-0.5]]
   pdf (MkGaussian mean cov) x ===# pure 0.016427375
 
-partial
 gaussianCDF : Device => Property
 gaussianCDF = fixedProperty $ do
   let gaussian = MkGaussian (tensor [[0.5]]) (tensor [[[1.44]]])
@@ -55,7 +52,7 @@ gaussianCDF = fixedProperty $ do
   cdf gaussian (tensor [[0.5]]) ===# pure 0.5
   cdf gaussian (tensor [[1.5]]) ===# pure 0.7976716
 
-export partial
+export
 group : Device => Group
 group = MkGroup "Distribution" $ [
       ("Gaussian univariate pdf", gaussianUnivariatePDF)
diff --git a/test/runner/Unit/TestTensor.idr b/test/runner/Unit/TestTensor.idr
index c0077fc0d..b2970a205 100644
--- a/test/runner/Unit/TestTensor.idr
+++ b/test/runner/Unit/TestTensor.idr
@@ -34,7 +34,6 @@ import Utils.Comparison
 import Utils.Cases
 import Utils.Proof
 
-partial
 tensorThenEval : Device => Property
 tensorThenEval @{device} = withTests 20 . property $ do
   shape <- forAll shapes
@@ -54,7 +53,6 @@ tensorThenEval @{device} = withTests 20 . property $ do
   x <- forAll (literal shape bool)
   x === unsafePerformIO (eval device $ pure $ tensor {dtype = PRED} x)
 
-partial
 evalTuple : Device => Property
 evalTuple @{device} = property $ do
   s0 <- forAll shapes
@@ -86,7 +84,6 @@ evalTuple @{device} = property $ do
   x1' === x1
   x2' === x2
 
-partial
 evalTupleNonTrivial : Device => Property
 evalTupleNonTrivial @{device} = property $ do
   let xs = do let y0 = tensor [1.0, -2.0, 0.4]
@@ -101,7 +98,6 @@ evalTupleNonTrivial @{device} = property $ do
   v ==~ Scalar (exp (-2.0) + 3.0)
   w ==~ [| exp [1.0, -2.0] |]
 
-partial
 canConvertAtXlaNumericBounds : Device => Property
 canConvertAtXlaNumericBounds = fixedProperty $ do
   let f64min : Literal [] Double = min @{Finite}
@@ -140,7 +136,6 @@ canConvertAtXlaNumericBounds = fixedProperty $ do
   unsafeEval (tensor u64min == min') === True
   unsafeEval (tensor u64max == max') === True
 
-partial
 boundedNonFinite : Device => Property
 boundedNonFinite = fixedProperty $ do
   let min' : Tensor [] S32 = Types.min @{NonFinite}
@@ -163,7 +158,6 @@ boundedNonFinite = fixedProperty $ do
   unsafeEval {dtype = F64} (Types.min @{NonFinite}) === -inf
   unsafeEval {dtype = F64} (Types.max @{NonFinite}) === inf
 
-partial
 iota : Device => Property
 iota = withTests 20 . property $ do
   init <- forAll shapes
@@ -190,7 +184,6 @@ iota = withTests 20 . property $ do
 
   actual ===# castDtype rangeFull
 
-partial
 iotaExamples : Device => Property
 iotaExamples = fixedProperty $ do
   iota 0 ===# tensor {dtype = S32} [0, 1, 2, 3]
@@ -212,7 +205,6 @@ iotaExamples = fixedProperty $ do
                                     [1.0, 1.0, 1.0, 1.0, 1.0],
                                     [2.0, 2.0, 2.0, 2.0, 2.0]]
 
-partial
 show : Device => Property
 show = fixedProperty $ do
   let x : Tag $ Tensor [] S32 = pure 1
@@ -281,7 +273,6 @@ show = fixedProperty $ do
     """
   -- x ===# pure 24  -- bug in XLA? https://github.com/openxla/xla/issues/14299
 
-partial
 cast : Device => Property
 cast = property $ do
   shape <- forAll shapes
@@ -298,7 +289,6 @@ cast = property $ do
   let x : Tensor shape F64 = castDtype $ tensor {dtype = S32} lit
   x ===# tensor (map (cast {to = Double}) lit)
 
-partial
 identity : Device => Property
 identity = fixedProperty $ do
   identity ===# tensor {dtype = S32} []
@@ -314,7 +304,7 @@ identity = fixedProperty $ do
     ]
 
 namespace Vector
-  export partial
+  export
   (@@) : Device => Property
   (@@) = fixedProperty $ do
     let l = tensor {dtype = S32} [-2, 0, 1]
@@ -322,7 +312,7 @@ namespace Vector
     l @@ r ===# -4
 
 namespace Matrix
-  export partial
+  export
   (@@) : Device => Property
   (@@) = fixedProperty $ do
     let l = tensor {dtype = S32} [[-2, 0, 1], [1, 3, 4]]
@@ -333,7 +323,6 @@ namespace Matrix
         r = tensor {dtype = S32} [[3, -1], [3, 2], [-1, -4]]
     l @@ r ===# tensor [[ -7,  -2], [  8, -11]]
 
-partial
 dotGeneral : Device => Property
 dotGeneral = fixedProperty $ do
   dotGeneral [] [] [] [] 2 3 ===# tensor {dtype = S32} 6
@@ -390,7 +379,6 @@ dotGeneral = fixedProperty $ do
                           [1.1037, 1.5626]]]
   dotGeneral [0] [0] [2] [1] l r ===# expected
 
-partial
 argmin : Device => Property
 argmin = property $ do
   d <- forAll dims
@@ -398,7 +386,6 @@ argmin = property $ do
   let xs = tensor xs
   (do pure $ slice [at !(argmin xs)] xs) ===# reduce [0] @{Min} xs
 
-partial
 argmax : Device => Property
 argmax = property $ do
   d <- forAll dims
@@ -406,7 +393,6 @@ argmax = property $ do
   let xs = tensor xs
   (do pure $ slice [at !(argmax xs)] xs) ===# reduce [0] @{Max} xs
 
-partial
 select : Device => Property
 select = fixedProperty $ do
   let onTrue = tensor {dtype = S32} 1
@@ -420,14 +406,12 @@ select = fixedProperty $ do
       expected = tensor [[6, 1, 2], [3, 10, 11]]
   select pred onTrue onFalse ===# expected
 
-partial
 erf : Device => Property
 erf = fixedProperty $ do
   let x = tensor [-1.5, -0.5, 0.5, 1.5]
       expected = tensor [-0.96610516, -0.5204998, 0.5204998, 0.9661051]
   erf x ===# expected
 
-partial
 cholesky : Device => Property
 cholesky = fixedProperty $ do
   let x = tensor [[1.0, 0.0], [2.0, 0.0]]
@@ -447,7 +431,6 @@ cholesky = fixedProperty $ do
             ]
   cholesky x ===# expected
 
-partial
 triangularSolveResultAndInverse : Device => Property
 triangularSolveResultAndInverse = fixedProperty $ do
   let a = tensor [
@@ -478,7 +461,6 @@ triangularSolveResultAndInverse = fixedProperty $ do
   actual ===# expected
   a.T @@ actual ===# b
 
-partial
 triangularSolveIgnoresOppositeElems : Device => Property
 triangularSolveIgnoresOppositeElems = fixedProperty $ do
   let a = tensor [[1.0, 2.0], [3.0, 4.0]]
@@ -489,12 +471,11 @@ triangularSolveIgnoresOppositeElems = fixedProperty $ do
   let aUpper = tensor [[1.0, 2.0], [0.0, 4.0]]
   a \| b ===# aUpper \| b
 
-partial
 trace : Device => Property
 trace = fixedProperty $
   trace (tensor {dtype = S32} [[-1, 5], [1, 4]]) ===# pure 3
 
-export partial
+export
 group : Device => Group
 group = MkGroup "Tensor" $ [
       ("eval . tensor", tensorThenEval)
diff --git a/test/runner/Unit/TestTensor/Elementwise.idr b/test/runner/Unit/TestTensor/Elementwise.idr
index 5dc704c12..7c9b69db5 100644
--- a/test/runner/Unit/TestTensor/Elementwise.idr
+++ b/test/runner/Unit/TestTensor/Elementwise.idr
@@ -26,7 +26,7 @@ import Utils.Comparison
 import Utils.Cases
 
 namespace S32
-  export partial
+  export
   testElementwiseUnary :
     Device =>
     (Int32 -> Int32) ->
@@ -39,7 +39,7 @@ namespace S32
     [| fInt x |] === unsafeEval (fTensor x')
 
 namespace F64
-  export partial
+  export
   testElementwiseUnary :
     Device =>
     (Double -> Double) ->
@@ -52,7 +52,7 @@ namespace F64
     [| fDouble x |] ==~ unsafeEval (fTensor x')
 
 namespace PRED
-  export partial
+  export
   testElementwiseUnary :
     Device =>
     (Bool -> Bool) ->
@@ -65,7 +65,7 @@ namespace PRED
     [| fBool x |] === unsafeEval (fTensor x')
 
 namespace S32
-  export partial
+  export
   testElementwiseBinary :
     Device =>
     (Int32 -> Int32 -> Int32) ->
@@ -79,21 +79,18 @@ namespace S32
         y' = tensor {dtype = S32} y
     [| fInt x y |] === unsafeEval (fTensor x' y')
 
-partial
 div : Device => Property
 div = fixedProperty $ do
   div (tensor {shape = [0]} []) [] ===# tensor []
   div (fill 9) [Scalar 1, Scalar 2, Scalar 3, Scalar 4, Scalar 5] ===# tensor [9, 4, 3, 2, 1]
   div (fill 1) [Scalar 1, Scalar 2, Scalar 3] ===# tensor [1, 0, 0]
 
-partial
 rem : Device => Property
 rem = fixedProperty $ do
   rem (tensor {shape = [0]} []) [] ===# tensor []
   rem (fill 9) [Scalar 1, Scalar 2, Scalar 3, Scalar 4, Scalar 5] ===# tensor [0, 1, 0, 1, 4]
   rem (fill 1) [Scalar 1, Scalar 2, Scalar 3] ===# tensor [0, 1, 1]
 
-partial
 divAndRemReconstructOriginal : Device => Property
 divAndRemReconstructOriginal = property $ do
   [x, y] <- forAll (np [nats, nats])
@@ -105,7 +102,7 @@ divAndRemReconstructOriginal = property $ do
   tensor denom * div numer denom + rem numer denom ===# numer
 
 namespace F64
-  export partial
+  export
   testElementwiseBinary :
     Device =>
     (Double -> Double -> Double) ->
@@ -120,7 +117,7 @@ namespace F64
     [| fDouble x y |] ==~ unsafeEval (fTensor x' y')
 
 namespace PRED
-  export partial
+  export
   testElementwiseBinary :
     Device =>
     (Bool -> Bool -> Bool) ->
@@ -134,7 +131,6 @@ namespace PRED
         y' = tensor {dtype = PRED} y
     [| fBool x y |] === unsafeEval (fTensor x' y')
 
-partial
 scalarMultiplication : Device => Property
 scalarMultiplication = property $ do
   shape <- forAll shapes
@@ -146,7 +142,6 @@ scalarMultiplication = property $ do
           scalar' = tensor {dtype = F64} (Scalar scalar)
       map (scalar *) lit ==~ unsafeEval (scalar' * lit')
 
-partial
 scalarDivision : Device => Property
 scalarDivision = property $ do
   shape <- forAll shapes
@@ -159,7 +154,7 @@ scalarDivision = property $ do
       map (/ scalar) lit ==~ unsafeEval (lit' / scalar')
 
 namespace S32
-  export partial
+  export
   testElementwiseComparator :
     Device =>
     (Int32 -> Int32 -> Bool) ->
@@ -174,7 +169,7 @@ namespace S32
     [| fInt x y |] === unsafeEval (fTensor x' y')
 
 namespace F64
-  export partial
+  export
   testElementwiseComparator :
     Device =>
     (Double -> Double -> Bool) ->
@@ -189,7 +184,7 @@ namespace F64
     [| fDouble x y |] === unsafeEval (fTensor x' y')
 
 namespace PRED
-  export partial
+  export
   testElementwiseComparator :
     Device =>
     (Bool -> Bool -> Bool) ->
@@ -197,7 +192,6 @@ namespace PRED
     Property
   testElementwiseComparator = testElementwiseBinary
 
-partial
 neutralIsNeutralForSum : Device => Property
 neutralIsNeutralForSum = property $ do
   shape <- forAll shapes
@@ -216,7 +210,6 @@ neutralIsNeutralForSum = property $ do
   unsafeEval right === x
   unsafeEval left === x
 
-partial
 neutralIsNeutralForProd : Device => Property
 neutralIsNeutralForProd = property $ do
   shape <- forAll shapes
@@ -235,7 +228,6 @@ neutralIsNeutralForProd = property $ do
   unsafeEval right === x
   unsafeEval left === x
 
-partial
 neutralIsNeutralForAny : Device => Property
 neutralIsNeutralForAny = property $ do
   shape <- forAll shapes
@@ -246,7 +238,6 @@ neutralIsNeutralForAny = property $ do
   unsafeEval right === x
   unsafeEval left === x
 
-partial
 neutralIsNeutralForAll : Device => Property
 neutralIsNeutralForAll = property $ do
   shape <- forAll shapes
@@ -257,7 +248,6 @@ neutralIsNeutralForAll = property $ do
   unsafeEval right === x
   unsafeEval left === x
 
-partial
 neutralIsNeutralForMin : Device => Property
 neutralIsNeutralForMin = property $ do
   shape <- forAll shapes
@@ -268,7 +258,6 @@ neutralIsNeutralForMin = property $ do
   unsafeEval right ==~ x
   unsafeEval left ==~ x
 
-partial
 neutralIsNeutralForMax : Device => Property
 neutralIsNeutralForMax = property $ do
   shape <- forAll shapes
@@ -306,7 +295,7 @@ acosh x = log (x + sqrt (x * x - 1))
 atanh : Double -> Double
 atanh x = log ((1 + x) / (1 - x)) / 2
 
-export partial
+export
 all : Device => List (PropertyName, Property)
 all = [
       ("negate S32", S32.testElementwiseUnary negate negate)
diff --git a/test/runner/Unit/TestTensor/HigherOrder.idr b/test/runner/Unit/TestTensor/HigherOrder.idr
index 1cae27021..ef5c56de2 100644
--- a/test/runner/Unit/TestTensor/HigherOrder.idr
+++ b/test/runner/Unit/TestTensor/HigherOrder.idr
@@ -25,7 +25,6 @@ import Utils
 import Utils.Comparison
 import Utils.Cases
 
-partial
 mapResult : Device => Property
 mapResult = property $ do
   shape <- forAll shapes
@@ -39,7 +38,6 @@ mapResult = property $ do
   let x' = tensor {dtype = S32} x
   map (+ 1) x === Tag.unsafeEval (map (pure . (+ 1)) x')
 
-partial
 mapNonTrivial : Device => Property
 mapNonTrivial = fixedProperty $ do
   let res : Tag (Tensor [] S32) = pure 2
@@ -55,7 +53,6 @@ mapNonTrivial = fixedProperty $ do
         ) (tensor 7)
   x ===# pure 28
 
-partial
 map2Result : Device => Property
 map2Result = fixedProperty $ do
   shape <- forAll shapes
@@ -73,13 +70,11 @@ map2Result = fixedProperty $ do
       y' = tensor {dtype = F64} y
   [| x + y |] ==~ unsafeEval (map2 (pure .: Tensor.(+)) x' y')
 
-partial
 map2ResultWithReusedFnArgs : Device => Property
 map2ResultWithReusedFnArgs = fixedProperty $ do
   let x : Tensor [] S32 = 6
   map2 (\x, y => pure $ x + x + y + y) 1 2 ===# pure x
 
-partial
 reduce : Device => Property
 reduce = fixedProperty $ do
   let x = tensor {dtype = S32} [[1, 2, 3], [-1, -2, -3]]
@@ -106,7 +101,6 @@ reduce = fixedProperty $ do
   let x = tensor {dtype = PRED} [[True, False, True], [True, False, False]]
   reduce @{All} [1] x ===# pure (tensor [False, False])
 
-partial
 sort : Device => Property
 sort = withTests 20 . property $ do
   d <- forAll dims
@@ -161,7 +155,6 @@ sort = withTests 20 . property $ do
   reflex : {n : _} -> LTE n n
   reflex = reflexive {ty = Nat}
 
-partial
 sortWithEmptyAxis : Device => Property
 sortWithEmptyAxis = fixedProperty $ do
   let x = tensor {shape = [0, 2, 3]} {dtype = S32} []
@@ -176,7 +169,6 @@ sortWithEmptyAxis = fixedProperty $ do
   let x = tensor {shape = [2, 0, 3]} {dtype = S32} [[], []]
   sort (<) 1 x ===# pure x
 
-partial
 sortWithRepeatedElements : Device => Property
 sortWithRepeatedElements = fixedProperty $ do
   let x = tensor {dtype = S32} [1, 3, 4, 3, 2]
@@ -186,7 +178,6 @@ sortWithRepeatedElements = fixedProperty $ do
   sort (<) 0 x ===# pure (tensor [[1, 2, 4], [3, 4, 5]])
   sort (<) 1 x ===# pure (tensor [[1, 4, 4], [2, 3, 5]])
 
-partial
 condResultTrivialUsage : Device => Property
 condResultTrivialUsage = fixedProperty $ do
   let x = tensor {dtype = S32} 0
@@ -204,7 +195,6 @@ condResultTrivialUsage = fixedProperty $ do
       y = tensor {dtype = S32} [[6, 7], [8, 9]]
   cond (tensor False) (\x => pure $ tensor 5 * x) x (pure . diag) y ===# pure (tensor [6, 9])
 
-partial
 condResultWithReusedArgs : Device => Property
 condResultWithReusedArgs = fixedProperty $ do
   let x = tensor {dtype = S32} 1
@@ -222,7 +212,7 @@ condResultWithReusedArgs = fixedProperty $ do
   cond (tensor True) (f (+)) x (f (*)) y ===# pure 2
   cond (tensor False) (f (+)) x (f (*)) y ===# pure 9
 
-export partial
+export
 all : Device => List (PropertyName, Property)
 all = [
       ("map", mapResult)
diff --git a/test/runner/Unit/TestTensor/Sampling.idr b/test/runner/Unit/TestTensor/Sampling.idr
index 32f33508a..08aad6927 100644
--- a/test/runner/Unit/TestTensor/Sampling.idr
+++ b/test/runner/Unit/TestTensor/Sampling.idr
@@ -31,7 +31,6 @@ range n = cast (Vect.range n)
 product1 : (x : Nat) -> product (the (List Nat) [x]) = x
 product1 x = rewrite plusZeroRightNeutral x in Refl
 
-partial
 iidKolmogorovSmirnov :
   {shape : _} -> Tensor shape F64 -> (Tensor shape F64 -> Tensor shape F64) -> Tag $ Tensor [] F64
 iidKolmogorovSmirnov samples cdf = do
@@ -47,7 +46,6 @@ iidKolmogorovSmirnov samples cdf = do
 Prelude.Ord a => Prelude.Ord (Literal [] a) where
   compare (Scalar x) (Scalar y) = compare x y
 
-partial
 uniform : Device => Property
 uniform = withTests 20 . property $ do
   bound <- forAll (literal [5] finiteDoubles)
@@ -70,7 +68,6 @@ uniform = withTests 20 . property $ do
 
   diff (unsafeEval ksTest) (<) 0.015
 
-partial
 uniformForNonFiniteBounds : Device => Property
 uniformForNonFiniteBounds = property $ do
   key <- forAll (literal [] nats)
@@ -88,7 +85,6 @@ uniformForNonFiniteBounds = property $ do
 minNormal : Literal [] Double
 minNormal = 2.23e-308  -- see https://en.wikipedia.org/wiki/IEEE_754
 
-partial
 uniformForFiniteEqualBounds : Device => Property
 uniformForFiniteEqualBounds = withTests 20 . property $ do
   key <- forAll (literal [] nats)
@@ -99,7 +95,6 @@ uniformForFiniteEqualBounds = withTests 20 . property $ do
 
   samples ===# pure bound
 
-partial
 uniformSeedIsUpdated : Device => Property
 uniformSeedIsUpdated = withTests 20 . property $ do
   bound <- forAll (literal [10] doubles)
@@ -122,7 +117,6 @@ uniformSeedIsUpdated = withTests 20 . property $ do
   diff seed'' (/=) seed'
   diff sample' (/=) sample
 
-partial
 uniformIsReproducible : Device => Property
 uniformIsReproducible = withTests 20 . property $ do
   bound <- forAll (literal [10] doubles)
@@ -143,7 +137,6 @@ uniformIsReproducible = withTests 20 . property $ do
 
   sample ==~ sample'
 
-partial
 normal : Device => Property
 normal = withTests 20 . property $ do
   key <- forAll (literal [] nats)
@@ -161,7 +154,6 @@ normal = withTests 20 . property $ do
 
   diff (unsafeEval ksTest) (<) 0.02
 
-partial
 normalSeedIsUpdated : Device => Property
 normalSeedIsUpdated = withTests 20 . property $ do
   key <- forAll (literal [] nats)
@@ -179,7 +171,6 @@ normalSeedIsUpdated = withTests 20 . property $ do
   diff seed'' (/=) seed'
   diff sample' (/=) sample
 
-partial
 normalIsReproducible : Device => Property
 normalIsReproducible = withTests 20 . property $ do
   key <- forAll (literal [] nats)
@@ -195,7 +186,7 @@ normalIsReproducible = withTests 20 . property $ do
 
   sample ==~ sample'
 
-export partial
+export
 all : Device => List (PropertyName, Property)
 all = [
       ("uniform", uniform)
diff --git a/test/runner/Unit/TestTensor/Slice.idr b/test/runner/Unit/TestTensor/Slice.idr
index 7b6d75788..7d86c519c 100644
--- a/test/runner/Unit/TestTensor/Slice.idr
+++ b/test/runner/Unit/TestTensor/Slice.idr
@@ -51,7 +51,6 @@ namespace MultiSlice
     slice {shape = [3, 4]} [0.to 2, at 2] === [2]
     slice {shape = [3, 4]} [at 1, at 2] === Prelude.Nil
 
-partial
 sliceStaticIndex : Device => Property
 sliceStaticIndex = fixedProperty $ do
   let x = tensor {dtype = S32} [3, 4, 5]
@@ -64,7 +63,6 @@ sliceStaticIndex = fixedProperty $ do
 
   slice [at idx] x ===# tensor 5
 
-partial
 sliceStaticSlice : Device => Property
 sliceStaticSlice = fixedProperty $ do
   let x = tensor {dtype = S32} [3, 4, 5]
@@ -78,7 +76,6 @@ sliceStaticSlice = fixedProperty $ do
   slice [2.to 2] x ===# tensor []
   slice [2.to 3] x ===# tensor [5]
 
-partial
 sliceStaticMixed : Device => Property
 sliceStaticMixed = fixedProperty $ do
   let x = tensor {dtype = S32} [[3, 4, 5], [6, 7, 8]]
@@ -100,7 +97,6 @@ sliceStaticMixed = fixedProperty $ do
 u64 : Nat -> Tensor [] U64
 u64 = tensor . Scalar
 
-partial
 sliceDynamicIndex : Device => Property
 sliceDynamicIndex = fixedProperty $ do
   let x = tensor {dtype = S32} [3, 4, 5]
@@ -116,7 +112,6 @@ sliceDynamicIndex = fixedProperty $ do
   slice [at (u64 2)] x ===# tensor [6, 7, 8]
   slice [at (u64 4)] x ===# tensor [6, 7, 8]
 
-partial
 sliceDynamicSlice : Device => Property
 sliceDynamicSlice = fixedProperty $ do
   let x = tensor {dtype = S32} [3, 4, 5]
@@ -145,7 +140,6 @@ sliceDynamicSlice = fixedProperty $ do
   slice [(u64 4).size 0] x ===# tensor []
   slice [(u64 4).size 1] x ===# tensor [[6, 7, 8]]
 
-partial
 sliceMixed : Device => Property
 sliceMixed = fixedProperty $ do
   let x = tensor {dtype = S32} [[3, 4, 5], [6, 7, 8]]
@@ -187,7 +181,6 @@ index : (idx : Nat) -> {auto 0 inDim : LT idx n} -> Literal [n] a -> Literal []
 index {inDim = (LTESucc _)} 0 (y :: _) = y
 index {inDim = (LTESucc _)} (S k) (_ :: xs) = index k xs
 
-partial
 sliceForVariableIndex : Device => Property
 sliceForVariableIndex = property $ do
   idx <- forAll dims
@@ -202,7 +195,7 @@ sliceForVariableIndex = property $ do
   inDim {idx = 0} = LTESucc LTEZero
   inDim {idx = (S k)} = LTESucc inDim
 
-export partial
+export
 all : Device => List (PropertyName, Property)
 all = [
       ("MultiSlice.slice", MultiSlice.slice)
diff --git a/test/runner/Unit/TestTensor/Structure.idr b/test/runner/Unit/TestTensor/Structure.idr
index de5c69404..a6cd46f25 100644
--- a/test/runner/Unit/TestTensor/Structure.idr
+++ b/test/runner/Unit/TestTensor/Structure.idr
@@ -27,7 +27,6 @@ import Utils
 import Utils.Comparison
 import Utils.Cases
 
-partial
 reshape : Device => Property
 reshape = fixedProperty $ do
   reshape 3 ===# tensor {dtype = S32} [3]
@@ -46,7 +45,6 @@ reshape = fixedProperty $ do
   let flattened = tensor {dtype = S32} [3, 4, 5, 6, 7, 8]
   reshape x ===# flattened
 
-partial
 expand : Device => Property
 expand = fixedProperty $ do
   expand 0 3 ===# tensor {dtype = S32} [3]
@@ -55,7 +53,6 @@ expand = fixedProperty $ do
       withExtraDim = tensor [[[3, 4, 5]], [[6, 7, 8]]]
   expand 1 x ===# withExtraDim
 
-partial
 broadcast : Device => Property
 broadcast = fixedProperty $ do
   broadcast {to = []} {dtype = S32} 7 ===# 7
@@ -105,7 +102,6 @@ broadcast = fixedProperty $ do
       ]
   broadcast {to = [2, 2, 5, 3]} x ===# expected
 
-partial
 triangle : Device => Property
 triangle = fixedProperty $ do
   let x = tensor {dtype = S32} []
@@ -124,7 +120,6 @@ triangle = fixedProperty $ do
   triangle Upper x ===# tensor [[1, 2, 3], [0, 5, 6], [0, 0, 9]]
   triangle Lower x ===# tensor [[1, 0, 0], [4, 5, 0], [7, 8, 9]]
 
-partial
 diag : Device => Property
 diag = fixedProperty $ do
   let x = tensor {dtype = S32} []
@@ -136,7 +131,6 @@ diag = fixedProperty $ do
   let x = tensor {dtype = S32} [[1, 2], [3, 4]]
   diag x ===# tensor [1, 4]
 
-partial
 concat : Device => Property
 concat = fixedProperty $ do
   let vector = tensor {dtype = S32} [3, 4, 5]
@@ -220,7 +214,6 @@ broadcastableCannotStackDimensionGtOne : Broadcastable [3, 2] [3, 7] -> Void
 broadcastableCannotStackDimensionGtOne (Match Same) impossible
 broadcastableCannotStackDimensionGtOne (Nest Same) impossible
 
-partial
 squeeze : Device => Property
 squeeze = fixedProperty $ do
   let x = tensor {dtype = S32} [[3]]
@@ -238,7 +231,6 @@ squeeze = fixedProperty $ do
 squeezableCannotRemoveNonOnes : Squeezable [1, 2] [] -> Void
 squeezableCannotRemoveNonOnes (Nest _) impossible
 
-partial
 (.T) : Device => Property
 (.T) = fixedProperty $ do
   (tensor {dtype = S32} []).T ===# tensor []
@@ -248,7 +240,6 @@ partial
       expected = tensor [[1, 4, 7], [2, 5, 8], [3, 6, 9]]
   x.T ===# expected
 
-partial
 transpose : Device => Property
 transpose = fixedProperty $ do
   let x = tensor {dtype = S32} [[0, 1], [2, 3]]
@@ -287,7 +278,6 @@ transpose = fixedProperty $ do
   slice [all, at 1, at 0] (transpose [0, 2, 1, 3] x) ===# slice [all, at 0, at 1] x
   slice [at 2, at 4, at 0, at 1] (transpose [2, 3, 1, 0] x) ===# slice [at 1, at 0, at 2, at 4] x
 
-partial
 reverse : Device => Property
 reverse = fixedProperty $ do
   let x = tensor {shape = [0]} {dtype = S32} []
@@ -315,7 +305,7 @@ reverse = fixedProperty $ do
     [[[ 1,  0], [ 3,  2]], [[ 5,  4], [ 7,  6]], [[ 9,  8], [11, 10]]]
   ]
 
-export partial
+export
 all : Device => List (PropertyName, Property)
 all = [
       ("reshape", reshape)
diff --git a/test/runner/Utils.idr b/test/runner/Utils.idr
index f15c03153..83e1679f2 100644
--- a/test/runner/Utils.idr
+++ b/test/runner/Utils.idr
@@ -29,20 +29,20 @@ isNan : Double -> Bool
 isNan x = x /= x
 
 namespace Tag
-  export partial
+  export
   unsafeEval : Device => PrimitiveRW dtype ty => Tag (Tensor shape dtype) -> Literal shape ty
   unsafeEval @{device} = unsafePerformIO . eval device
 
-export partial
+export
 unsafeEval : Device => PrimitiveRW dtype ty => Tensor shape dtype -> Literal shape ty
 unsafeEval @{device} x = unsafePerformIO $ eval device (pure x)
 
 namespace TensorList
   namespace Tag
-    export partial
+    export
     unsafeEval : Device => Tag (TensorList shapes tys) -> All2 Literal shapes tys
     unsafeEval @{device} = unsafePerformIO . eval device
 
-  export partial
+  export
   unsafeEval : Device => TensorList shapes tys -> All2 Literal shapes tys
   unsafeEval @{device} xs = unsafePerformIO $ eval device (pure xs)
diff --git a/test/runner/Utils/Comparison.idr b/test/runner/Utils/Comparison.idr
index 696eb6f82..d2501cc0e 100644
--- a/test/runner/Utils/Comparison.idr
+++ b/test/runner/Utils/Comparison.idr
@@ -57,7 +57,7 @@ export infix 6 ===#
 
 namespace Tag
   namespace PRED
-    export partial
+    export
     (===#) : Device =>
              Monad m =>
              {shape : _} ->
@@ -67,7 +67,7 @@ namespace Tag
     x ===# y = unsafeEval x === unsafeEval y
 
   namespace S32
-    export partial
+    export
     (===#) : Device =>
              Monad m =>
              {shape : _} ->
@@ -77,7 +77,7 @@ namespace Tag
     x ===# y = unsafeEval x === unsafeEval y
 
   namespace U32
-    export partial
+    export
     (===#) : Device =>
              Monad m =>
              {shape : _} ->
@@ -87,7 +87,7 @@ namespace Tag
     x ===# y = unsafeEval x === unsafeEval y
 
   namespace U64
-    export partial
+    export
     (===#) : Device =>
              Monad m =>
              {shape : _} ->
@@ -97,7 +97,7 @@ namespace Tag
     x ===# y = unsafeEval x === unsafeEval y
 
   namespace F64
-    export partial
+    export
     (===#) : Device =>
              Monad m =>
              {shape : _} ->
@@ -108,27 +108,27 @@ namespace Tag
     x ===# y = (==~) {tol} (unsafeEval x) (unsafeEval y)
 
 namespace PRED
-  export partial
+  export
   (===#) : Device => Monad m => {shape : _} -> Tensor shape PRED -> Tensor shape PRED -> TestT m ()
   x ===# y = unsafeEval (pure x) === unsafeEval (pure y)
 
 namespace S32
-  export partial
+  export
   (===#) : Device => Monad m => {shape : _} -> Tensor shape S32 -> Tensor shape S32 -> TestT m ()
   x ===# y = unsafeEval (pure x) === unsafeEval (pure y)
 
 namespace U32
-  export partial
+  export
   (===#) : Device => Monad m => {shape : _} -> Tensor shape U32 -> Tensor shape U32 -> TestT m ()
   x ===# y = unsafeEval (pure x) === unsafeEval (pure y)
 
 namespace U64
-  export partial
+  export
   (===#) : Device => Monad m => {shape : _} -> Tensor shape U64 -> Tensor shape U64 -> TestT m ()
   x ===# y = unsafeEval (pure x) === unsafeEval (pure y)
 
 namespace F64
-  export partial
+  export
   (===#) : Device => Monad m => {shape : _} -> {default floatingPointTolerance tol : Double} ->
            Tensor shape F64 -> Tensor shape F64 -> TestT m ()
   x ===# y = (==~) {tol} (unsafeEval $ pure x) (unsafeEval $ pure y)