// Copyright 2018 The TensorFlow Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

#if !COMPILING_TENSORFLOW_MODULE
import TensorFlow
#endif

@differentiable
public func meanSquaredError<Scalar: FloatingPoint>(predicted: Tensor<Scalar>,
                                                    expected: Tensor<Scalar>) -> Tensor<Scalar> {
    return (expected - predicted).squared().mean()
}

// Copyright 2018 The TensorFlow Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

#if !COMPILING_TENSORFLOW_MODULE
import TensorFlow
#endif

@differentiable(vjp: _vjpMSE)
public func meanSquaredError<Scalar: FloatingPoint>(predicted: Tensor<Scalar>,
                                                    expected: Tensor<Scalar>) -> Tensor<Scalar> {
    return (expected - predicted).squared().mean()
}

@usableFromInline func _vjpMSE<Scalar: FloatingPoint>(
    predicted: Tensor<Scalar>, expected: Tensor<Scalar>
) -> (Tensor<Scalar>, (Tensor<Scalar>) -> (Tensor<Scalar>, Tensor<Scalar>)) {
    let result = meanSquaredError(predicted: predicted, expected: expected)
    return (result, { [n = predicted.scalarCountTensor] v in
      let d = 2 * v / Tensor(n)
      return (d, -d)
    })
}

--- 
+++ 
@@ -16,18 +16,8 @@
 import TensorFlow
 #endif
 
-@differentiable(vjp: _vjpMSE)
+@differentiable
 public func meanSquaredError<Scalar: FloatingPoint>(predicted: Tensor<Scalar>,
                                                     expected: Tensor<Scalar>) -> Tensor<Scalar> {
     return (expected - predicted).squared().mean()
 }
-
-@usableFromInline func _vjpMSE<Scalar: FloatingPoint>(
-    predicted: Tensor<Scalar>, expected: Tensor<Scalar>
-) -> (Tensor<Scalar>, (Tensor<Scalar>) -> (Tensor<Scalar>, Tensor<Scalar>)) {
-    let result = meanSquaredError(predicted: predicted, expected: expected)
-    return (result, { [n = predicted.scalarCountTensor] v in
-      let d = 2 * v / Tensor(n)
-      return (d, -d)
-    })
-}
Use generic AD on all functions.
