Review

Author: Shashi456

Sources/TensorFlow/Operators/Math.swift

@@ -649,16 +649,16 @@ public func rsqrt<T: TensorFlowFloatingPoint>(_ x: Tensor<T>) -> Tensor<T> {
 public func cosineSimilarity<Scalar: TensorFlowFloatingPoint>(
     _ x: Tensor<Scalar>, _ y: Tensor<Scalar>
 ) -> Tensor<Scalar> {
-    return (x * y).sum() /
-        (sqrt(x.squared().sum()) * sqrt(y.squared().sum()))
+    (x * y).sum() / (sqrt(x.squared().sum()) * sqrt(y.squared().sum()))
 }
 
-/// Returns the cosine distance between `x` and `y`.
-@differentiable(wrt: (x, y))
+/// Returns the cosine distance between `x` and `y`. Cosine distance is defined as
+/// `1 - cosineSimilarity(x, y)`.
+@differentiable
 public func cosineDistance<Scalar: TensorFlowFloatingPoint>(
     _ x: Tensor<Scalar>, _ y: Tensor<Scalar>
 ) -> Tensor<Scalar> {
-    return 1 - cosineSimilarity(x, y)
+    1 - cosineSimilarity(x, y)
 }
 
 @inlinable