 
     /// Returns the mean and variance of this tensor along the specified axes. The reduced
     /// dimensions are retained with value `1`.
+    ///
+    /// - Parameter axes: The dimensions to reduce.
+    /// - Precondition: Each value in `axes` must be in the range `-rank..<rank`.
+    @inlinable
+    @differentiable(wrt: self)
+    func moments(alongAxes axes: [Int]) -> Moments<Scalar> {
+        // TODO(TF-433): Remove workaround for differentiating `map`.
+        let axes = {axes.map(Int32.init)}()
+        return moments(alongAxes: Tensor<Int32>(axes))
+    }
+
+    /// Returns the mean and variance of this tensor along the specified axes. The reduced
+    /// dimensions are retained with value `1`.
+    ///
+    /// - Parameter axes: The dimensions to reduce.
+    /// - Precondition: Each value in `axes` must be in the range `-rank..<rank`.
     @inlinable
     @differentiable(wrt: self)
     func moments(alongAxes axes: Int...) -> Moments<Scalar> {

public extension Tensor where Scalar: TensorFlowFloatingPoint {
    /// Returns the mean and variance of this tensor along the specified axes. The reduced
    /// dimensions are removed.
    @inlinable
    @differentiable(wrt: self)
    func moments(squeezingAxes axes: [Int]) -> Moments<Scalar> {
        let mean = self.mean(alongAxes: axes)
        let variance = squaredDifference(self, mean).mean(alongAxes: axes)
        return Moments<Scalar>(
            mean: mean.squeezingShape(at: axes),
            variance: variance.squeezingShape(at: axes))
    }

    /// Returns the mean and variance of this tensor along the specified axes. The reduced
    /// dimensions are removed.
    @inlinable
    @differentiable(wrt: self)
    func moments(squeezingAxes axes: Int...) -> Moments<Scalar> {


    /// Returns the mean and variance of this tensor along the specified axes. The reduced
    /// dimensions are retained with value `1`.
    ///
    /// - Parameter axes: The dimensions to reduce.
    /// - Precondition: `axes` must have rank `1`.
    /// - Precondition: Each value in `axes` must be in the range `-rank..<rank`.
    @inlinable
    @differentiable(wrt: self)
    func moments(alongAxes axes: Tensor<Int32>) -> Moments<Scalar> {
        let mean = self.mean(alongAxes: axes)
        let variance = squaredDifference(self, mean).mean(alongAxes: axes)
        return Moments<Scalar>(mean: mean, variance: variance)
