Moved reduceMean to Ops

tensorflow-mnist/app/Main.hs

@@ -41,9 +41,6 @@ randomParam width (TF.Shape shape) =
   where
     stddev = TF.scalar (1 / sqrt (fromIntegral width))
 
-reduceMean :: TF.Tensor TF.Build Float -> TF.Tensor TF.Build Float
-reduceMean xs = TF.mean xs (TF.scalar (0 :: Int32))
-
 -- Types must match due to model structure.
 type LabelType = Int32
 
@@ -85,12 +82,12 @@ createModel = do
     labels <- TF.placeholder [batchSize]
     let labelVecs = TF.oneHot labels (fromIntegral numLabels) 1 0
         loss =
-            reduceMean $ fst $ TF.softmaxCrossEntropyWithLogits logits labelVecs
+            TF.reduceMean $ fst $ TF.softmaxCrossEntropyWithLogits logits labelVecs
         params = [hiddenWeights, hiddenBiases, logitWeights, logitBiases]
     trainStep <- TF.minimizeWith TF.adam loss params
 
     let correctPredictions = TF.equal predict labels
-    errorRateTensor <- TF.render $ 1 - reduceMean (TF.cast correctPredictions)
+    errorRateTensor <- TF.render $ 1 - TF.reduceMean (TF.cast correctPredictions)
 
     return Model {
           train = \imFeed lFeed -> TF.runWithFeeds_ [

tensorflow-ops/src/TensorFlow/Ops.hs

@@ -106,6 +106,8 @@ module TensorFlow.Ops
     , CoreOps.range
     , CoreOps.range'
     , reducedShape
+    , reduceMean
+    , reduceMean'
     , CoreOps.relu
     , CoreOps.relu'
     , CoreOps.reluGrad
@@ -330,6 +332,23 @@ reduceSum' :: (OneOf '[ Double, Float, Int32, Int64
 reduceSum' params x = CoreOps.sum' params x allAxes
   where allAxes = CoreOps.range 0 (CoreOps.rank x :: Tensor Build Int32) 1
 
+-- | Computes the mean of elements across dimensions of a tensor.
+-- See `TensorFlow.GenOps.Core.mean`
+reduceMean
+  :: ( TensorType a
+     , OneOf '[ Double, Float, Complex Float, Complex Double] a
+     )
+  => Tensor v a -> Tensor Build a
+reduceMean = reduceMean' id
+
+reduceMean'
+  :: ( TensorType a
+     , OneOf '[ Double, Float, Complex Float, Complex Double] a
+     )
+  => OpParams -> Tensor v a -> Tensor Build a
+reduceMean' params x = CoreOps.mean' params x allAxes
+  where allAxes = CoreOps.range 0 (CoreOps.rank x :: Tensor Build Int32) 1
+
 -- | Create a constant vector.
 vector :: TensorType a => [a] -> Tensor Build a
 vector = vector' id

tensorflow-ops/tests/MatrixTest.hs

@@ -6,7 +6,7 @@ import Control.Monad (replicateM_)
 
 import qualified Data.Vector as V
 import qualified TensorFlow.Core as TF
-import qualified TensorFlow.GenOps.Core as TF (square, rank)
+import qualified TensorFlow.GenOps.Core as TF (square)
 import qualified TensorFlow.Minimize as TF
 import qualified TensorFlow.Ops as TF hiding (initializedVariable)
 import qualified TensorFlow.Variable as TF
@@ -18,17 +18,14 @@ import TensorFlow.Test (assertAllClose)
 randomParam :: TF.Shape -> TF.Session (TF.Tensor TF.Value Float)
 randomParam (TF.Shape shape) = TF.truncatedNormal (TF.vector shape)
 
-reduceMean :: TF.Tensor v Float -> TF.Tensor TF.Build Float
-reduceMean xs = TF.mean xs (TF.range 0 (TF.rank xs) 1)
-
 fitMatrix :: Test
 fitMatrix = testCase "fitMatrix" $ TF.runSession $ do
   u <- TF.initializedVariable =<< randomParam [2, 1]
   v <- TF.initializedVariable =<< randomParam [1, 2]
   let ones = [1, 1, 1, 1] :: [Float]
       matx = TF.constant [2, 2] ones
       diff = matx `TF.sub` (TF.readValue u `TF.matMul` TF.readValue v)
-      loss = reduceMean $ TF.square diff
+      loss = TF.reduceMean $ TF.square diff
   trainStep <- TF.minimizeWith (TF.gradientDescent 0.01) loss [u, v]
   replicateM_ 1000 (TF.run trainStep)
   (u',v') <- TF.run (TF.readValue u, TF.readValue v)