Added support for 'Tensor.cumulativeProduct(alongAxis:exclusive:reverse:)'. (#358)

Author: eaplatanios

Sources/TensorFlow/Operators/Math.swift

@@ -1963,6 +1963,79 @@ public extension Tensor where Scalar: Numeric {
     ) -> Tensor {
         Raw.cumsum(self, axis: axis, exclusive: exclusive, reverse: reverse)
     }
+
+    /// Returns the cumulative product of this tensor along the specified axis. By default, this
+    /// function performs an inclusive cumulative product which means that the first element of the
+    /// input is identical to the first element of the output:
+    /// ```
+    /// Tensor<Float>([a, b, c]).cumulativeProduct() = Tensor<Float>([a, a * b, a * b * c])
+    /// ```
+    /// By setting the `exclusive` argument to `true`, an exclusive cumulative product is performed
+    /// instead:
+    /// ```
+    /// Tensor<Float>([a, b, c]).cumulativeProduct(exclusive: true) = Tensor<Float>([1, a, a * b])
+    /// ```
+    /// By setting the `reverse` argument to `true`, the cumulative product is performed in the
+    /// opposite direction:
+    /// ```
+    /// Tensor<Float>([a, b, c]).cumulativeProduct(reverse: true) ==
+    ///     Tensor<Float>([a * b * c, a * b, a])
+    /// ```
+    /// This is more efficient than separately reversing the resulting tensor.
+    ///
+    /// - Parameters:
+    ///   - axis: Axis along which to perform the cumulative product operation.
+    ///   - exclusive: Indicates whether to perform an exclusive cumulative product.
+    ///   - reverse: Indicates whether to perform the cumulative product in reversed order.
+    /// - Returns: Result of the cumulative product operation.
+    /// - Precondition: `axis` must be in the range `-rank..<rank`.
+    @inlinable
+    @differentiable(wrt: self where Scalar: TensorFlowFloatingPoint)
+    func cumulativeProduct(
+        alongAxis axis: Int,
+        exclusive: Bool = false,
+        reverse: Bool = false
+    ) -> Tensor {
+        cumulativeProduct(
+            alongAxis: Tensor<Int32>(Int32(axis)),
+            exclusive: exclusive,
+            reverse: reverse)
+    }
+
+    /// Returns the cumulative product of this tensor along the specified axis. By default, this
+    /// function performs an inclusive cumulative product which means that the first element of the
+    /// input is identical to the first element of the output:
+    /// ```
+    /// Tensor<Float>([a, b, c]).cumulativeProduct() = Tensor<Float>([a, a * b, a * b * c])
+    /// ```
+    /// By setting the `exclusive` argument to `true`, an exclusive cumulative product is performed
+    /// instead:
+    /// ```
+    /// Tensor<Float>([a, b, c]).cumulativeProduct(exclusive: true) = Tensor<Float>([1, a, a * b])
+    /// ```
+    /// By setting the `reverse` argument to `true`, the cumulative product is performed in the
+    /// opposite direction:
+    /// ```
+    /// Tensor<Float>([a, b, c]).cumulativeProduct(reverse: true) ==
+    ///     Tensor<Float>([a * b * c, a * b, a])
+    /// ```
+    /// This is more efficient than separately reversing the resulting tensor.
+    ///
+    /// - Parameters:
+    ///   - axis: Axis along which to perform the cumulative product operation.
+    ///   - exclusive: Indicates whether to perform an exclusive cumulative product.
+    ///   - reverse: Indicates whether to perform the cumulative product in reversed order.
+    /// - Returns: Result of the cumulative product operation.
+    /// - Precondition: `axis` must be in the range `-rank..<rank`.
+    @inlinable
+    @differentiable(wrt: self, vjp: _vjpCumulativeProduct where Scalar: TensorFlowFloatingPoint)
+    func cumulativeProduct(
+        alongAxis axis: Tensor<Int32>,
+        exclusive: Bool = false,
+        reverse: Bool = false
+    ) -> Tensor {
+        Raw.cumprod(self, axis: axis, exclusive: exclusive, reverse: reverse)
+    }
 }
 
 internal extension Tensor where Scalar: TensorFlowFloatingPoint {
@@ -2008,6 +2081,22 @@ internal extension Tensor where Scalar: TensorFlowFloatingPoint {
             v.cumulativeSum(alongAxis: axis, exclusive: exclusive, reverse: !reverse)
         })
     }
+
+    @inlinable
+    func _vjpCumulativeProduct(
+        alongAxis axis: Tensor<Int32>,
+        exclusive: Bool = false,
+        reverse: Bool = false
+    ) -> (Tensor, (Tensor) -> Tensor) {
+        let result = cumulativeProduct(alongAxis: axis, exclusive: exclusive, reverse: reverse)
+        return (result, { v in
+            (result * v).cumulativeSum(
+                alongAxis: axis,
+                exclusive: exclusive,
+                reverse: !reverse
+            ) / self
+        })
+    }
 }
 
 // TODO: Consider making the return type be generic over `FloatingPoint` types

Tests/TensorFlowTests/OperatorTests/MathTests.swift

@@ -242,6 +242,33 @@ final class MathOperatorTests: XCTestCase {
         XCTAssertEqual(reverseExclusiveCumsum1, Tensor<Float>([[3, 2, 0], [9, 5, 0]]))
     }
 
+    func testCumulativeProduct() {
+        // 2 x 3
+        let x = Tensor<Float>([[0, 1, 2], [3, 4, 5]])
+        let cumprod0 = x.cumulativeProduct(alongAxis: 0)
+        let cumprod1 = x.cumulativeProduct(alongAxis: 1)
+        let exclusiveCumprod0 = x.cumulativeProduct(alongAxis: 0, exclusive: true)
+        let exclusiveCumprod1 = x.cumulativeProduct(alongAxis: 1, exclusive: true)
+        let reverseCumprod0 = x.cumulativeProduct(alongAxis: 0, reverse: true)
+        let reverseCumprod1 = x.cumulativeProduct(alongAxis: 1, reverse: true)
+        let reverseExclusiveCumprod0 = x.cumulativeProduct(
+            alongAxis: 0,
+            exclusive: true,
+            reverse: true)
+        let reverseExclusiveCumprod1 = x.cumulativeProduct(
+            alongAxis: 1,
+            exclusive: true,
+            reverse: true)
+        XCTAssertEqual(cumprod0, Tensor<Float>([[0, 1, 2], [0, 4, 10]]))
+        XCTAssertEqual(cumprod1, Tensor<Float>([[0, 0, 0], [3, 12, 60]]))
+        XCTAssertEqual(exclusiveCumprod0, Tensor<Float>([[1, 1, 1], [0, 1, 2]]))
+        XCTAssertEqual(exclusiveCumprod1, Tensor<Float>([[1, 0, 0], [1, 3, 12]]))
+        XCTAssertEqual(reverseCumprod0, Tensor<Float>([[0, 4, 10], [3, 4, 5]]))
+        XCTAssertEqual(reverseCumprod1, Tensor<Float>([[0, 2, 2], [60, 20, 5]]))
+        XCTAssertEqual(reverseExclusiveCumprod0, Tensor<Float>([[3, 4, 5], [1, 1, 1]]))
+        XCTAssertEqual(reverseExclusiveCumprod1, Tensor<Float>([[2, 2, 1], [20, 5, 1]]))
+    }
+
     func testStandardDeviation() {
         XCTAssertEqual(Tensor<Float>([1]).standardDeviation(), Tensor(0))
         XCTAssertEqual(Tensor<Float>([0, 1]).standardDeviation(alongAxes: 0), Tensor(0.5))
@@ -485,6 +512,7 @@ final class MathOperatorTests: XCTestCase {
         ("testArgmax", testArgmax),
         ("testReduction", testReduction),
         ("testCumulativeSum", testCumulativeSum),
+        ("testCumulativeProduct", testCumulativeProduct),
         ("testStandardDeviation", testStandardDeviation),
         ("testLogSumExp", testLogSumExp),
         ("testMoments", testMoments),