<commit_before>// Copyright 2019 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.

infix operator .!=: ComparisonPrecedence

/// Returns a tensor with the same shape and scalars as the specified tensor.
@inlinable
@differentiable(where Scalar: TensorFlowFloatingPoint)
public func identity<Scalar>(_ x: Tensor<Scalar>) -> Tensor<Scalar> {
    return x
}

//===------------------------------------------------------------------------------------------===//
// Shape Transformations
//===------------------------------------------------------------------------------------------===//

public extension TensorFlowScalar {
    /// Convert to a tensor with the specified rank, with all dimensions equal to `1`.
    @inlinable
    func makeTensor(rank: Int) -> Tensor<Self> {
        return Tensor(repeating: self, shape: TensorShape(rank))
    }
}

public extension Tensor {
    /// Unpacks the given dimension of a rank-`R` tensor into multiple rank-`(R-1)` tensors. Unpacks
    /// `N` tensors from this tensor by chipping it along the `axis` dimension, where `N` is
    /// inferred from this tensor's shape. For example, given a tensor with shape `[A, B, C, D]`:
    ///
    ///   - If `axis == 0` then the `i`-th tensor in the returned array is the slice
    ///     `self[i, :, :, :]` and each tensor in that array will have shape `[B, C, D]`.
    ///     (Note that the dimension unpacked along is gone, unlike
    ///     `Tensor.split(numSplits:alongAxis)`, or `Tensor.split(sizes:alongAxis)`).
    ///   - If `axis == 1` then the `i`-th tensor in the returned array is the slice
    ///     `value[:, i, :, :]` and each tensor in that array will have shape `[A, C, D]`.
    ///   - Etc.
    ///
    /// This is the opposite of `Tensor.init(stacking:alongAxis:)`.
    ///
    /// - Parameters:
    ///   - axis: Dimension along which to unstack. Negative values wrap around.
    ///
    /// - Precondition: `axis` must be in the range `[-rank, rank)`, where `rank` is the rank of the
    ///   provided tensors.
    ///
    /// - Returns: Array containing the unstacked tensors.
    @inlinable
    @differentiable(vjp: _vjpUnstacked(alongAxis:) where Scalar: TensorFlowFloatingPoint)
    func unstacked(alongAxis axis: Int = 0) -> [Tensor] {
        return Raw.unpack(value: self, num: Int64(shape[axis]), axis: Int64(axis))
    }

    /// Splits a tensor into multiple tensors. The tensor is split along dimension `axis` into
    /// `count` smaller tensors. This requires that `count` evenly divides `shape[axis]`.
    ///
    /// For example:
    /// ```
    /// // 'value' is a tensor with shape [5, 30]
    /// // Split 'value' into 3 tensors along dimension 1:
    /// let parts = value.split(count: 3, alongAxis: 1)
    /// parts[0] // has shape [5, 10]
    /// parts[1] // has shape [5, 10]
    /// parts[2] // has shape [5, 10]
    /// ```
    ///
    /// - Parameters:
    ///   - count: Number of splits to create.
    ///   - axis: The dimension along which to split this tensor. Negative values wrap around.
    ///
    /// - Precondition: `count` must divide the size of dimension `axis` evenly.
    /// - Precondition: `axis` must be in the range `[-rank, rank)`, where `rank` is the rank of the
    ///   provided tensors.
    ///
    /// - Returns: An array containing the tensors parts.
    @inlinable
    @differentiable(vjp: _vjpSplit(count:alongAxis:) where Scalar: TensorFlowFloatingPoint)
    func split(count: Int, alongAxis axis: Int = 0) -> [Tensor] {
        return Raw.split(
            splitDim: Tensor<Int32>(Int32(axis)), value: self, numSplit: Int64(count))
    }

    /// Splits a tensor into multiple tensors. The tensor is split  into `sizes.shape[0]` pieces.
    /// The shape of the `i`-th piece has the same shape as this tensor except along dimension
    /// `axis` where the size is `sizes[i]`.
    ///
    /// For example:
    /// ```
    /// // 'value' is a tensor with shape [5, 30]
    /// // Split 'value' into 3 tensors with sizes [4, 15, 11] along dimension 1:
    /// let parts = value.split(sizes: Tensor<Int32>([4, 15, 11]), alongAxis: 1)
    /// parts[0] // has shape [5, 4]
    /// parts[1] // has shape [5, 15]
    /// parts[2] // has shape [5, 11]
    /// ```
    ///
    /// - Parameters:
    ///   - sizes: 1-D tensor containing the size of each split.
    ///   - axis: Dimension along which to split this tensor. Negative values wrap around.
    ///
    /// - Precondition: The values in `sizes` must add up to the size of dimension `axis`.
    /// - Precondition: `axis` must be in the range `[-rank, rank)`, where `rank` is the rank of the
    ///   provided tensors.
    ///
    /// - Returns: Array containing the tensors parts.
    @inlinable
    @differentiable(
        wrt: self,
        vjp: _vjpSplit(sizes:alongAxis:) where Scalar: TensorFlowFloatingPoint)
    func split(sizes: Tensor<Int32>, alongAxis axis: Int = 0) -> [Tensor] {
        return Raw.splitV(
            value: self,
            sizeSplits: sizes,
            splitDim: Tensor<Int32>(Int32(axis)),
            numSplit: Int64(sizes.shape[0]))
    }

    /// Returns a tiled tensor, constructed by tiling this tensor.
    ///
    /// This constructor creates a new tensor by replicating this tensor `multiples` times. The
    /// constructed tensor's `i`'th dimension has `self.shape[i] * multiples[i]` elements, and the
    /// values of this tensor are replicated `multiples[i]` times along the `i`'th dimension. For
    /// example, tiling `[a b c d]` by `[2]` produces `[a b c d a b c d]`.
    ///
    /// - Precondition: The shape of `multiples` must be `[tensor.rank]`.
    @inlinable
    @differentiable(wrt: self, vjp: _vjpTiled(multiples:) where Scalar: TensorFlowFloatingPoint)
    func tiled(multiples: Tensor<Int32>) -> Tensor {
        return Raw.tile(self, multiples: multiples)
    }

    /// Reshape to the shape of the specified `Tensor`.
    /// - Precondition: The number of scalars matches the new shape.
    @inlinable
    @differentiable(wrt: self where Scalar: TensorFlowFloatingPoint)
    func reshaped<T>(like other: Tensor<T>) -> Tensor {
        return reshaped(toShape: other.shapeTensor)
    }

    /// Reshape to the specified shape.
    /// - Precondition: The number of scalars matches the new shape.
    @inlinable
    @differentiable(wrt: self where Scalar: TensorFlowFloatingPoint)
    func reshaped(to newShape: TensorShape) -> Tensor {
        // TODO(TF-433): Remove workaround for differentiating `map`.
        return reshaped(toShape: Tensor<Int32>({newShape.dimensions.map(Int32.init)}()))
    }

    /// Reshape to the specified `Tensor` representing a shape.
    /// - Precondition: The number of scalars matches the new shape.
    @inlinable
    @differentiable(
        wrt: self,
        vjp: _vjpReshaped(toShape:) where Scalar: TensorFlowFloatingPoint)
    func reshaped(toShape newShape: Tensor<Int32>) -> Tensor {
        return Raw.reshape(self, shape: newShape)
    }

    /// Return a copy of the tensor collapsed into a 1-D `Tensor`, in row-major order.
    @inlinable
    @differentiable(wrt: self where Scalar: TensorFlowFloatingPoint)
    func flattened() -> Tensor {
        return reshaped(to: [-1])
    }

    /// Returns a shape-expanded `Tensor`, with a dimension of 1 inserted at the specified shape
    /// indices.
    @inlinable
    @differentiable(wrt: self where Scalar: TensorFlowFloatingPoint)
    func expandingShape(at axes: Int...) -> Tensor {
        return expandingShape(at: axes)
    }

    /// Returns a shape-expanded `Tensor`, with a dimension of 1 inserted at the
    /// specified shape indices.
    @inlinable
    @differentiable(wrt: self, vjp: _vjpExpandingShape(at:) where Scalar: TensorFlowFloatingPoint)
    func expandingShape(at axes: [Int]) -> Tensor {
        var result = self
        for i in axes { result = Raw.expandDims(result, dim: Tensor<Int32>(Int32(i))) }
        return result
    }

    /// Returns a rank-lifted `Tensor` with a leading dimension of 1.
    @inlinable
    @differentiable(wrt: self where Scalar: TensorFlowFloatingPoint)
    func rankLifted() -> Tensor {
        return expandingShape(at: 0)
    }

    /// Removes the specified dimensions of size 1 from the shape of a tensor. If no dimensions are
    /// specified, then all dimensions of size 1 will be removed.
    @inlinable
    @differentiable(wrt: self where Scalar: TensorFlowFloatingPoint)
    func squeezingShape(at axes: Int...) -> Tensor {
        return squeezingShape(at: axes)
    }

    /// Removes the specified dimensions of size 1 from the shape of a tensor. If no dimensions are
    /// specified, then all dimensions of size 1 will be removed.
    @inlinable
    @differentiable(wrt: self, vjp: _vjpSqueezingShape(at:) where Scalar: TensorFlowFloatingPoint)
    func squeezingShape(at axes: [Int]) -> Tensor {
        return Raw.squeeze(self, squeezeDims: axes.map(Int32.init))
    }
}

internal extension Tensor where Scalar: TensorFlowFloatingPoint {
    @inlinable
    func _vjpUnstacked(
        alongAxis axis: Int = 0
    ) -> ([Tensor], (Array<Tensor>.TangentVector) -> Tensor) {
        let result = unstacked(alongAxis: axis)
        return (result, { v in Tensor(stacking: v.base, alongAxis: axis) })
    }

    @inlinable
    func _vjpTiled(
        multiples: Tensor<Int32>
    ) -> (Tensor, (Tensor) -> Tensor) {
        return (tiled(multiples: multiples), { [shape = shapeTensor] v in
            let splitShape = Tensor<Int32>(stacking: [multiples, shape]).transposed().flattened()
            let axes = Tensor<Int32>(rangeFrom: 0, to: Int32(splitShape.scalarCount), stride: 2)
            return v.reshaped(toShape: splitShape).sum(squeezingAxes: axes)
        })
    }

    @inlinable
    func _vjpSplit(
        count: Int,
        alongAxis axis: Int = 0
    ) -> ([Tensor], (Array<Tensor>.TangentVector) -> Tensor) {
        let result = split(count: count, alongAxis: axis)
        return (result, { v in Tensor(concatenating: v.base, alongAxis: axis) })
    }

    @inlinable
    func _vjpSplit(
        sizes: Tensor<Int32>,
        alongAxis axis: Int = 0
    ) -> ([Tensor], (Array<Tensor>.TangentVector) -> Tensor) {
        let result = split(sizes: sizes, alongAxis: axis)
        return (result, { v in Tensor(concatenating: v.base, alongAxis: axis) })
    }

    @inlinable
    func _vjpReshaped(toShape newShape: Tensor<Int32>) -> (Tensor, (Tensor) -> Tensor) {
        let value = reshaped(toShape: newShape)
        return (value, { [shape = shapeTensor] v in v.reshaped(toShape: shape) })
    }

    @inlinable
    func _vjpExpandingShape(at axes: [Int]) -> (Tensor, (Tensor) -> Tensor) {
        let value = self.expandingShape(at: axes)
        return (value, { v in v.squeezingShape(at: axes) })
    }

    @inlinable
    func _vjpSqueezingShape(at axes: [Int]) -> (Tensor, (Tensor) -> Tensor) {
        let value = squeezingShape(at: axes)
        return (value, { [shape = shapeTensor] v in v.reshaped(toShape: shape) })
    }
}

//===------------------------------------------------------------------------------------------===//
// Other Tensor Transformations
//===------------------------------------------------------------------------------------------===//

infix operator ++: AdditionPrecedence

public extension Tensor {
    /// Returns a transposed tensor, with dimensions permuted in the specified order.
    @inlinable
    @differentiable(
        wrt: self,
        vjp: _vjpTransposed(withPermutations:) where Scalar: TensorFlowFloatingPoint)
    func transposed(withPermutations permutations: Tensor<Int32>) -> Tensor {
        return Raw.transpose(self, perm: permutations)
    }

    /// Returns a transposed tensor, with dimensions permuted in the specified order.
    @inlinable
    @differentiable(
        wrt: self,
        vjp: _vjpTransposed(withPermutations:) where Scalar: TensorFlowFloatingPoint)
    func transposed(withPermutations permutations: [Int]) -> Tensor {
        let permutations = permutations.map(Int32.init)
        return transposed(withPermutations: Tensor<Int32>(permutations))
    }

    /// Returns a transposed tensor, with dimensions permuted in the specified order.
    @inlinable
    @differentiable(
        wrt: self, vjp: _vjpTransposed(withPermutations:) where Scalar: TensorFlowFloatingPoint)
    func transposed(withPermutations permutations: Int...) -> Tensor {
        return transposed(withPermutations: permutations)
    }

    /// Returns a transposed tensor, with dimensions permuted in reverse order.
    @inlinable
    @differentiable(wrt: self, vjp: _vjpTransposed() where Scalar: TensorFlowFloatingPoint)
    func transposed() -> Tensor {
        let defaultPermutations = rankTensor - 1 - Tensor<Int32>(
            rangeFrom: 0, to: Int32(rank), stride: 1)
        return transposed(withPermutations: Tensor<Int32>(defaultPermutations))
    }

    /// Returns a concatenated tensor along the specified axis.
    /// - Precondition: The tensors must have the same dimensions, except for the
    ///   specified axis.
    /// - Precondition: The axis must be in the range `-rank..<rank`.
    @inlinable
    @differentiable(vjp: _vjpConcatenated where Scalar: TensorFlowFloatingPoint)
    func concatenated(with other: Tensor, alongAxis axis: Int = 0) -> Tensor {
        return Tensor(concatenating: [self, other], alongAxis: axis)
    }

    /// Concatenation operator.
    /// - Note: `++` is a custom operator that does not exist in Swift, but does
    ///   in Haskell/Scala. Its addition is not an insignificant language change
    ///   and may be controversial. The existence/naming of `++` will be discussed
    ///   during a later API design phase.
    @inlinable
    @differentiable(where Scalar: TensorFlowFloatingPoint)
    static func ++ (lhs: Tensor, rhs: Tensor) -> Tensor {
        return lhs.concatenated(with: rhs)
    }

    /// Returns a tensor by gathering slices of the input at `indices` along the `axis` dimension
    ///
    /// For 0-D (scalar) `indices`:
    /// ```
    /// result[p_0,          ..., p_{axis-1},
    ///        p_{axis + 1}, ..., p_{N-1}] =
    /// self[p_0,          ..., p_{axis-1},
    ///      indices,
    ///      p_{axis + 1}, ..., p_{N-1}]
    /// ```
    ///
    /// For 1-D (vector) `indices`:
    /// ```
    /// result[p_0,          ..., p_{axis-1},
    ///        i,
    ///        p_{axis + 1}, ..., p_{N-1}] =
    /// self[p_0,          ..., p_{axis-1},
    ///      indices[i],
    ///      p_{axis + 1}, ..., p_{N-1}]
    /// ```
    ///
    /// In the general case, produces a resulting tensor where:
    /// ```
    /// result[p_0,             ..., p_{axis-1},
    ///        i_{batch\_dims}, ..., i_{M-1},
    ///        p_{axis + 1},    ..., p_{N-1}] =
    /// self[p_0,             ..., p_{axis-1},
    ///      indices[i_0,     ..., i_{M-1}],
    ///      p_{axis + 1},    ..., p_{N-1}]
    /// ```
    /// where `N = self.rank` and `M = indices.rank`.
    ///
    /// The shape of the resulting tensor is:
    /// `self.shape[..<axis] + indices.shape + self.shape[(axis + 1)...]`.
    ///
    /// - Note: On CPU, if an out-of-range index is found, an error is thrown. On GPU, if an
    /// out-of-range index is found, a 0 is stored in the corresponding output values.
    ///
    /// - Parameters:
    ///   - indices: Contains the indices to gather at.
    ///   - axis: Dimension along which to gather. Negative values wrap around.
    ///
    /// - Precondition: `axis` must be in the range `[-rank, rank)`.
    ///
    /// - Returns: The gathered tensor.
    @inlinable
    @differentiable(wrt: self, vjp: _vjpGathering where Scalar : TensorFlowFloatingPoint)
    func gathering(atIndices indices: Tensor<Int32>, alongAxis axis: Int = 0) -> Tensor {
        return Raw.gatherV2(params: self, indices: indices, axis: Tensor<Int32>(Int32(axis)))
    }

    /// Returns slices of this tensor at `indices`, while ignoring the first `batchDims` dimensions
    /// that correspond to batch dimensions. The gather is performed along the first non-batch
    /// dimension.
    ///
    /// Performs similar functionality to `gathering`, except that the resulting tensor shape is 
    /// now:
    /// ```
    /// self.shape[..<batchDims] + 
    ///   indices.shape[batchDims...] + 
    ///   self.shape[(batchDims + indices.rank + 1)...]
    /// ```
    ///
    /// - Parameters:
    ///   - indices: Contains the indices to gather.
    ///   - batchDims: Number of leading batch dimensions to ignore.
    ///
    /// - Precondition: `batchDims` must be less than `indices.rank`.
    ///
    /// - Returns: The gathered tensor.
    @inlinable
    @differentiable(wrt: self where Scalar: TensorFlowFloatingPoint)
    func batchGathering(atIndices indices: Tensor<Int32>) -> Tensor {
        var batchIndices = indices
        var accumulated = Tensor<Int32>(ones: [])
        accumulated *= Swift.withoutDerivative(at: shapeTensor) { $0[1] }
        let dValue = Swift.withoutDerivative(at: shapeTensor) { $0[0] }
        let dIndices = Tensor<Int32>(
            rangeFrom: Tensor<Int32>(zeros: []),
            to: dValue,
            stride: Tensor<Int32>(ones: [])
        ) * accumulated
        let dShape = Tensor<Int32>(concatenating: [
            dValue.rankLifted(),
            Tensor<Int32>([Int32](repeating: 1, count: indices.rank - 1))])
        batchIndices += dIndices.reshaped(toShape: dShape)
        let flatIndices = batchIndices.flattened()
        let outerShape = Swift.withoutDerivative(at: shapeTensor) { $0[2...] }
        let innerShape = Swift.withoutDerivative(at: shapeTensor) { $0[..<2] }.product(squeezingAxes: [0])
        let flatTensor = reshaped(toShape: innerShape.rankLifted().concatenated(with: outerShape))
        let flatResult = flatTensor.gathering(atIndices: flatIndices)
        return flatResult.reshaped(toShape: indices.shapeTensor.concatenated(with: outerShape))
    }

    /// Returns a tensor by gathering the values after applying the provided boolean mask to the input.
    ///
    /// For example:
    /// ```
    /// // 1-D example
    /// // tensor is [0, 1, 2, 3]
    /// // mask is [true, false, true, false]
    /// tensor.gathering(where: mask) // is [0, 2]
    ///
    /// // 2-D example
    /// // tensor is [[1, 2], [3, 4], [5, 6]]
    /// // mask is [true, false, true]
    /// tensor.gathering(where: mask) // is [[1, 2], [5, 6]]
    /// ```
    ///
    /// In general, `0 < mask.rank = K <= tensor.rank`, and the `mask`'s shape must match the first
    /// K dimensions of the `tensor`'s shape. We then have:
    /// `tensor.gathering(where: mask)[i, j1, ..., jd] = tensor[i1, ..., iK, j1, ..., jd]`, where
    /// `[i1, ..., iK]` is the `i`th `true` entry of `mask` (row-major order).
    ///
    /// The `axis` could be used with `mask` to indicate the axis to mask from. In that case,
    /// `axis + mask.rank <= tensor.rank` and the `mask``'s shape must match the first
    /// `axis + mask.rank` dimensions of the `tensor`'s shape.
    ///
    /// - Parameters:
    ///   - mask: K-D boolean tensor, where `K <= self.rank`.
    ///   - axis: 0-D integer tensor representing the axis in `self` to mask from, where
    ///     `K + axis <= self.rank`.
    ///
    /// - Precondition: The `mask` cannot be a scalar: `mask.rank != 0`.
    ///
    /// - Returns: `(self.rank - K + 1)`-dimensional tensor populated by entries in this tensor
    ///   corresponding to `true` values in `mask`.
    @inlinable
    // @differentiable(wrt: self where Scalar: TensorFlowFloatingPoint)
    func gathering(where mask: Tensor<Bool>, alongAxis axis: Int = 0) -> Tensor {
        precondition(mask.rank != 0, "The boolean mask cannot be a scalar.")
        // TODO: Remove once control flow AD is supported.
        let rank = self.rank
        let posAxis = { axis < 0 ? axis + rank : axis }()
        let leadingSize = shapeTensor[posAxis ..< posAxis + mask.rank].product().rankLifted()
        let reshapedTensor = reshaped(
            toShape: Tensor<Int32>(concatenating: [
                shapeTensor[..<posAxis],
                leadingSize,<commit_msg>Made 'Tensor.gathering(where:alongAxis:)' differentiable. (#271)<commit_after>// Copyright 2019 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.

infix operator .!=: ComparisonPrecedence

/// Returns a tensor with the same shape and scalars as the specified tensor.
@inlinable
@differentiable(where Scalar: TensorFlowFloatingPoint)
public func identity<Scalar>(_ x: Tensor<Scalar>) -> Tensor<Scalar> {
    return x
}

//===------------------------------------------------------------------------------------------===//
// Shape Transformations
//===------------------------------------------------------------------------------------------===//

public extension TensorFlowScalar {
    /// Convert to a tensor with the specified rank, with all dimensions equal to `1`.
    @inlinable
    func makeTensor(rank: Int) -> Tensor<Self> {
        return Tensor(repeating: self, shape: TensorShape(rank))
    }
}

public extension Tensor {
    /// Unpacks the given dimension of a rank-`R` tensor into multiple rank-`(R-1)` tensors. Unpacks
    /// `N` tensors from this tensor by chipping it along the `axis` dimension, where `N` is
    /// inferred from this tensor's shape. For example, given a tensor with shape `[A, B, C, D]`:
    ///
    ///   - If `axis == 0` then the `i`-th tensor in the returned array is the slice
    ///     `self[i, :, :, :]` and each tensor in that array will have shape `[B, C, D]`.
    ///     (Note that the dimension unpacked along is gone, unlike
    ///     `Tensor.split(numSplits:alongAxis)`, or `Tensor.split(sizes:alongAxis)`).
    ///   - If `axis == 1` then the `i`-th tensor in the returned array is the slice
    ///     `value[:, i, :, :]` and each tensor in that array will have shape `[A, C, D]`.
    ///   - Etc.
    ///
    /// This is the opposite of `Tensor.init(stacking:alongAxis:)`.
    ///
    /// - Parameters:
    ///   - axis: Dimension along which to unstack. Negative values wrap around.
    ///
    /// - Precondition: `axis` must be in the range `[-rank, rank)`, where `rank` is the rank of the
    ///   provided tensors.
    ///
    /// - Returns: Array containing the unstacked tensors.
    @inlinable
    @differentiable(vjp: _vjpUnstacked(alongAxis:) where Scalar: TensorFlowFloatingPoint)
    func unstacked(alongAxis axis: Int = 0) -> [Tensor] {
        return Raw.unpack(value: self, num: Int64(shape[axis]), axis: Int64(axis))
    }

    /// Splits a tensor into multiple tensors. The tensor is split along dimension `axis` into
    /// `count` smaller tensors. This requires that `count` evenly divides `shape[axis]`.
    ///
    /// For example:
    /// ```
    /// // 'value' is a tensor with shape [5, 30]
    /// // Split 'value' into 3 tensors along dimension 1:
    /// let parts = value.split(count: 3, alongAxis: 1)
    /// parts[0] // has shape [5, 10]
    /// parts[1] // has shape [5, 10]
    /// parts[2] // has shape [5, 10]
    /// ```
    ///
    /// - Parameters:
    ///   - count: Number of splits to create.
    ///   - axis: The dimension along which to split this tensor. Negative values wrap around.
    ///
    /// - Precondition: `count` must divide the size of dimension `axis` evenly.
    /// - Precondition: `axis` must be in the range `[-rank, rank)`, where `rank` is the rank of the
    ///   provided tensors.
    ///
    /// - Returns: An array containing the tensors parts.
    @inlinable
    @differentiable(vjp: _vjpSplit(count:alongAxis:) where Scalar: TensorFlowFloatingPoint)
    func split(count: Int, alongAxis axis: Int = 0) -> [Tensor] {
        return Raw.split(
            splitDim: Tensor<Int32>(Int32(axis)), value: self, numSplit: Int64(count))
    }

    /// Splits a tensor into multiple tensors. The tensor is split  into `sizes.shape[0]` pieces.
    /// The shape of the `i`-th piece has the same shape as this tensor except along dimension
    /// `axis` where the size is `sizes[i]`.
    ///
    /// For example:
    /// ```
    /// // 'value' is a tensor with shape [5, 30]
    /// // Split 'value' into 3 tensors with sizes [4, 15, 11] along dimension 1:
    /// let parts = value.split(sizes: Tensor<Int32>([4, 15, 11]), alongAxis: 1)
    /// parts[0] // has shape [5, 4]
    /// parts[1] // has shape [5, 15]
    /// parts[2] // has shape [5, 11]
    /// ```
    ///
    /// - Parameters:
    ///   - sizes: 1-D tensor containing the size of each split.
    ///   - axis: Dimension along which to split this tensor. Negative values wrap around.
    ///
    /// - Precondition: The values in `sizes` must add up to the size of dimension `axis`.
    /// - Precondition: `axis` must be in the range `[-rank, rank)`, where `rank` is the rank of the
    ///   provided tensors.
    ///
    /// - Returns: Array containing the tensors parts.
    @inlinable
    @differentiable(
        wrt: self,
        vjp: _vjpSplit(sizes:alongAxis:) where Scalar: TensorFlowFloatingPoint)
    func split(sizes: Tensor<Int32>, alongAxis axis: Int = 0) -> [Tensor] {
        return Raw.splitV(
            value: self,
            sizeSplits: sizes,
            splitDim: Tensor<Int32>(Int32(axis)),
            numSplit: Int64(sizes.shape[0]))
    }

    /// Returns a tiled tensor, constructed by tiling this tensor.
    ///
    /// This constructor creates a new tensor by replicating this tensor `multiples` times. The
    /// constructed tensor's `i`'th dimension has `self.shape[i] * multiples[i]` elements, and the
    /// values of this tensor are replicated `multiples[i]` times along the `i`'th dimension. For
    /// example, tiling `[a b c d]` by `[2]` produces `[a b c d a b c d]`.
    ///
    /// - Precondition: The shape of `multiples` must be `[tensor.rank]`.
    @inlinable
    @differentiable(wrt: self, vjp: _vjpTiled(multiples:) where Scalar: TensorFlowFloatingPoint)
    func tiled(multiples: Tensor<Int32>) -> Tensor {
        return Raw.tile(self, multiples: multiples)
    }

    /// Reshape to the shape of the specified `Tensor`.
    /// - Precondition: The number of scalars matches the new shape.
    @inlinable
    @differentiable(wrt: self where Scalar: TensorFlowFloatingPoint)
    func reshaped<T>(like other: Tensor<T>) -> Tensor {
        return reshaped(toShape: other.shapeTensor)
    }

    /// Reshape to the specified shape.
    /// - Precondition: The number of scalars matches the new shape.
    @inlinable
    @differentiable(wrt: self where Scalar: TensorFlowFloatingPoint)
    func reshaped(to newShape: TensorShape) -> Tensor {
        // TODO(TF-433): Remove workaround for differentiating `map`.
        return reshaped(toShape: Tensor<Int32>({newShape.dimensions.map(Int32.init)}()))
    }

    /// Reshape to the specified `Tensor` representing a shape.
    /// - Precondition: The number of scalars matches the new shape.
    @inlinable
    @differentiable(
        wrt: self,
        vjp: _vjpReshaped(toShape:) where Scalar: TensorFlowFloatingPoint)
    func reshaped(toShape newShape: Tensor<Int32>) -> Tensor {
        return Raw.reshape(self, shape: newShape)
    }

    /// Return a copy of the tensor collapsed into a 1-D `Tensor`, in row-major order.
    @inlinable
    @differentiable(wrt: self where Scalar: TensorFlowFloatingPoint)
    func flattened() -> Tensor {
        return reshaped(to: [-1])
    }

    /// Returns a shape-expanded `Tensor`, with a dimension of 1 inserted at the specified shape
    /// indices.
    @inlinable
    @differentiable(wrt: self where Scalar: TensorFlowFloatingPoint)
    func expandingShape(at axes: Int...) -> Tensor {
        return expandingShape(at: axes)
    }

    /// Returns a shape-expanded `Tensor`, with a dimension of 1 inserted at the
    /// specified shape indices.
    @inlinable
    @differentiable(wrt: self, vjp: _vjpExpandingShape(at:) where Scalar: TensorFlowFloatingPoint)
    func expandingShape(at axes: [Int]) -> Tensor {
        var result = self
        for i in axes { result = Raw.expandDims(result, dim: Tensor<Int32>(Int32(i))) }
        return result
    }

    /// Returns a rank-lifted `Tensor` with a leading dimension of 1.
    @inlinable
    @differentiable(wrt: self where Scalar: TensorFlowFloatingPoint)
    func rankLifted() -> Tensor {
        return expandingShape(at: 0)
    }

    /// Removes the specified dimensions of size 1 from the shape of a tensor. If no dimensions are
    /// specified, then all dimensions of size 1 will be removed.
    @inlinable
    @differentiable(wrt: self where Scalar: TensorFlowFloatingPoint)
    func squeezingShape(at axes: Int...) -> Tensor {
        return squeezingShape(at: axes)
    }

    /// Removes the specified dimensions of size 1 from the shape of a tensor. If no dimensions are
    /// specified, then all dimensions of size 1 will be removed.
    @inlinable
    @differentiable(wrt: self, vjp: _vjpSqueezingShape(at:) where Scalar: TensorFlowFloatingPoint)
    func squeezingShape(at axes: [Int]) -> Tensor {
        return Raw.squeeze(self, squeezeDims: axes.map(Int32.init))
    }
}

internal extension Tensor where Scalar: TensorFlowFloatingPoint {
    @inlinable
    func _vjpUnstacked(
        alongAxis axis: Int = 0
    ) -> ([Tensor], (Array<Tensor>.TangentVector) -> Tensor) {
        let result = unstacked(alongAxis: axis)
        return (result, { v in Tensor(stacking: v.base, alongAxis: axis) })
    }

    @inlinable
    func _vjpTiled(
        multiples: Tensor<Int32>
    ) -> (Tensor, (Tensor) -> Tensor) {
        return (tiled(multiples: multiples), { [shape = shapeTensor] v in
            let splitShape = Tensor<Int32>(stacking: [multiples, shape]).transposed().flattened()
            let axes = Tensor<Int32>(rangeFrom: 0, to: Int32(splitShape.scalarCount), stride: 2)
            return v.reshaped(toShape: splitShape).sum(squeezingAxes: axes)
        })
    }

    @inlinable
    func _vjpSplit(
        count: Int,
        alongAxis axis: Int = 0
    ) -> ([Tensor], (Array<Tensor>.TangentVector) -> Tensor) {
        let result = split(count: count, alongAxis: axis)
        return (result, { v in Tensor(concatenating: v.base, alongAxis: axis) })
    }

    @inlinable
    func _vjpSplit(
        sizes: Tensor<Int32>,
        alongAxis axis: Int = 0
    ) -> ([Tensor], (Array<Tensor>.TangentVector) -> Tensor) {
        let result = split(sizes: sizes, alongAxis: axis)
        return (result, { v in Tensor(concatenating: v.base, alongAxis: axis) })
    }

    @inlinable
    func _vjpReshaped(toShape newShape: Tensor<Int32>) -> (Tensor, (Tensor) -> Tensor) {
        let value = reshaped(toShape: newShape)
        return (value, { [shape = shapeTensor] v in v.reshaped(toShape: shape) })
    }

    @inlinable
    func _vjpExpandingShape(at axes: [Int]) -> (Tensor, (Tensor) -> Tensor) {
        let value = self.expandingShape(at: axes)
        return (value, { v in v.squeezingShape(at: axes) })
    }

    @inlinable
    func _vjpSqueezingShape(at axes: [Int]) -> (Tensor, (Tensor) -> Tensor) {
        let value = squeezingShape(at: axes)
        return (value, { [shape = shapeTensor] v in v.reshaped(toShape: shape) })
    }
}

//===------------------------------------------------------------------------------------------===//
// Other Tensor Transformations
//===------------------------------------------------------------------------------------------===//

infix operator ++: AdditionPrecedence

public extension Tensor {
    /// Returns a transposed tensor, with dimensions permuted in the specified order.
    @inlinable
    @differentiable(
        wrt: self,
        vjp: _vjpTransposed(withPermutations:) where Scalar: TensorFlowFloatingPoint)
    func transposed(withPermutations permutations: Tensor<Int32>) -> Tensor {
        return Raw.transpose(self, perm: permutations)
    }

    /// Returns a transposed tensor, with dimensions permuted in the specified order.
    @inlinable
    @differentiable(
        wrt: self,
        vjp: _vjpTransposed(withPermutations:) where Scalar: TensorFlowFloatingPoint)
    func transposed(withPermutations permutations: [Int]) -> Tensor {
        let permutations = permutations.map(Int32.init)
        return transposed(withPermutations: Tensor<Int32>(permutations))
    }

    /// Returns a transposed tensor, with dimensions permuted in the specified order.
    @inlinable
    @differentiable(
        wrt: self, vjp: _vjpTransposed(withPermutations:) where Scalar: TensorFlowFloatingPoint)
    func transposed(withPermutations permutations: Int...) -> Tensor {
        return transposed(withPermutations: permutations)
    }

    /// Returns a transposed tensor, with dimensions permuted in reverse order.
    @inlinable
    @differentiable(wrt: self, vjp: _vjpTransposed() where Scalar: TensorFlowFloatingPoint)
    func transposed() -> Tensor {
        let defaultPermutations = rankTensor - 1 - Tensor<Int32>(
            rangeFrom: 0, to: Int32(rank), stride: 1)
        return transposed(withPermutations: Tensor<Int32>(defaultPermutations))
    }

    /// Returns a concatenated tensor along the specified axis.
    /// - Precondition: The tensors must have the same dimensions, except for the
    ///   specified axis.
    /// - Precondition: The axis must be in the range `-rank..<rank`.
    @inlinable
    @differentiable(vjp: _vjpConcatenated where Scalar: TensorFlowFloatingPoint)
    func concatenated(with other: Tensor, alongAxis axis: Int = 0) -> Tensor {
        return Tensor(concatenating: [self, other], alongAxis: axis)
    }

    /// Concatenation operator.
    /// - Note: `++` is a custom operator that does not exist in Swift, but does
    ///   in Haskell/Scala. Its addition is not an insignificant language change
    ///   and may be controversial. The existence/naming of `++` will be discussed
    ///   during a later API design phase.
    @inlinable
    @differentiable(where Scalar: TensorFlowFloatingPoint)
    static func ++ (lhs: Tensor, rhs: Tensor) -> Tensor {
        return lhs.concatenated(with: rhs)
    }

    /// Returns a tensor by gathering slices of the input at `indices` along the `axis` dimension
    ///
    /// For 0-D (scalar) `indices`:
    /// ```
    /// result[p_0,          ..., p_{axis-1},
    ///        p_{axis + 1}, ..., p_{N-1}] =
    /// self[p_0,          ..., p_{axis-1},
    ///      indices,
    ///      p_{axis + 1}, ..., p_{N-1}]
    /// ```
    ///
    /// For 1-D (vector) `indices`:
    /// ```
    /// result[p_0,          ..., p_{axis-1},
    ///        i,
    ///        p_{axis + 1}, ..., p_{N-1}] =
    /// self[p_0,          ..., p_{axis-1},
    ///      indices[i],
    ///      p_{axis + 1}, ..., p_{N-1}]
    /// ```
    ///
    /// In the general case, produces a resulting tensor where:
    /// ```
    /// result[p_0,             ..., p_{axis-1},
    ///        i_{batch\_dims}, ..., i_{M-1},
    ///        p_{axis + 1},    ..., p_{N-1}] =
    /// self[p_0,             ..., p_{axis-1},
    ///      indices[i_0,     ..., i_{M-1}],
    ///      p_{axis + 1},    ..., p_{N-1}]
    /// ```
    /// where `N = self.rank` and `M = indices.rank`.
    ///
    /// The shape of the resulting tensor is:
    /// `self.shape[..<axis] + indices.shape + self.shape[(axis + 1)...]`.
    ///
    /// - Note: On CPU, if an out-of-range index is found, an error is thrown. On GPU, if an
    /// out-of-range index is found, a 0 is stored in the corresponding output values.
    ///
    /// - Parameters:
    ///   - indices: Contains the indices to gather at.
    ///   - axis: Dimension along which to gather. Negative values wrap around.
    ///
    /// - Precondition: `axis` must be in the range `[-rank, rank)`.
    ///
    /// - Returns: The gathered tensor.
    @inlinable
    @differentiable(wrt: self, vjp: _vjpGathering where Scalar : TensorFlowFloatingPoint)
    func gathering(atIndices indices: Tensor<Int32>, alongAxis axis: Int = 0) -> Tensor {
        return Raw.gatherV2(params: self, indices: indices, axis: Tensor<Int32>(Int32(axis)))
    }

    /// Returns slices of this tensor at `indices`, while ignoring the first `batchDims` dimensions
    /// that correspond to batch dimensions. The gather is performed along the first non-batch
    /// dimension.
    ///
    /// Performs similar functionality to `gathering`, except that the resulting tensor shape is 
    /// now:
    /// ```
    /// self.shape[..<batchDims] + 
    ///   indices.shape[batchDims...] + 
    ///   self.shape[(batchDims + indices.rank + 1)...]
    /// ```
    ///
    /// - Parameters:
    ///   - indices: Contains the indices to gather.
    ///   - batchDims: Number of leading batch dimensions to ignore.
    ///
    /// - Precondition: `batchDims` must be less than `indices.rank`.
    ///
    /// - Returns: The gathered tensor.
    @inlinable
    @differentiable(wrt: self where Scalar: TensorFlowFloatingPoint)
    func batchGathering(atIndices indices: Tensor<Int32>) -> Tensor {
        var batchIndices = indices
        var accumulated = Tensor<Int32>(ones: [])
        accumulated *= Swift.withoutDerivative(at: shapeTensor) { $0[1] }
        let dValue = Swift.withoutDerivative(at: shapeTensor) { $0[0] }
        let dIndices = Tensor<Int32>(
            rangeFrom: Tensor<Int32>(zeros: []),
            to: dValue,
            stride: Tensor<Int32>(ones: [])
        ) * accumulated
        let dShape = Tensor<Int32>(concatenating: [
            dValue.rankLifted(),
            Tensor<Int32>([Int32](repeating: 1, count: indices.rank - 1))])
        batchIndices += dIndices.reshaped(toShape: dShape)
        let flatIndices = batchIndices.flattened()
        let outerShape = Swift.withoutDerivative(at: shapeTensor) { $0[2...] }
        let innerShape = Swift.withoutDerivative(at: shapeTensor) { $0[..<2] }.product(squeezingAxes: [0])
        let flatTensor = reshaped(toShape: innerShape.rankLifted().concatenated(with: outerShape))
        let flatResult = flatTensor.gathering(atIndices: flatIndices)
        return flatResult.reshaped(toShape: indices.shapeTensor.concatenated(with: outerShape))
    }

    /// Returns a tensor by gathering the values after applying the provided boolean mask to the input.
    ///
    /// For example:
    /// ```
    /// // 1-D example
    /// // tensor is [0, 1, 2, 3]
    /// // mask is [true, false, true, false]
    /// tensor.gathering(where: mask) // is [0, 2]
    ///
    /// // 2-D example
    /// // tensor is [[1, 2], [3, 4], [5, 6]]
    /// // mask is [true, false, true]
    /// tensor.gathering(where: mask) // is [[1, 2], [5, 6]]
    /// ```
    ///
    /// In general, `0 < mask.rank = K <= tensor.rank`, and the `mask`'s shape must match the first
    /// K dimensions of the `tensor`'s shape. We then have:
    /// `tensor.gathering(where: mask)[i, j1, ..., jd] = tensor[i1, ..., iK, j1, ..., jd]`, where
    /// `[i1, ..., iK]` is the `i`th `true` entry of `mask` (row-major order).
    ///
    /// The `axis` could be used with `mask` to indicate the axis to mask from. In that case,
    /// `axis + mask.rank <= tensor.rank` and the `mask``'s shape must match the first
    /// `axis + mask.rank` dimensions of the `tensor`'s shape.
    ///
    /// - Parameters:
    ///   - mask: K-D boolean tensor, where `K <= self.rank`.
    ///   - axis: 0-D integer tensor representing the axis in `self` to mask from, where
    ///     `K + axis <= self.rank`.
    ///
    /// - Precondition: The `mask` cannot be a scalar: `mask.rank != 0`.
    ///
    /// - Returns: `(self.rank - K + 1)`-dimensional tensor populated by entries in this tensor
    ///   corresponding to `true` values in `mask`.
    @inlinable
    @differentiable(wrt: self where Scalar: TensorFlowFloatingPoint)
    func gathering(where mask: Tensor<Bool>, alongAxis axis: Int = 0) -> Tensor {
        precondition(mask.rank != 0, "The boolean mask cannot be a scalar.")
        let posAxis = Swift.withoutDerivative(at: self.rank) { r in axis < 0 ? axis + r : axis }
        let leadingSize = shapeTensor[posAxis ..< posAxis + mask.rank].product().rankLifted()
        let reshapedTensor = reshaped(
            toShape: Tensor<Int32>(concatenating: [
                shapeTensor[..<posAxis],
                leadingSize,
{'fraction_non_alphanumeric': 0.15465338394951278, 'fraction_numerical': 0.011823875039292019, 'mean_word_length': 2.2766598003485976, 'pattern_counts': {'":': 0, '<': 101, '<?xml version=': 0, '>': 153, 'https://': 0, 'lorem ipsum': 0, 'www.': 2, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 2, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '6946746', 'n_tokens_mistral': 12949, 'n_tokens_neox': 12200, 'n_words': 3965}