<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.

import TensorFlow
import Python

struct TimeDistributed: Layer {
    var dense: Dense<Float>

    init(_ wrapped: Dense<Float>) {
        self.dense = wrapped
    }

    @differentiable(wrt: (self, input))
    func callAsFunction(_ input: Tensor<Float>) -> Tensor<Float> {
        let (batchSize, timeSteps, features) = (input.shape[0], input.shape[1], input.shape[2])
        let reshaped = input.reshaped(to: [batchSize * timeSteps, features])
        let output = dense(reshaped)
        let outputFeatures = output.shape[1]
        return output.reshaped(to: [batchSize, timeSteps, outputFeatures])
    }
}

struct FeedForward: Layer {
    var dense1: TimeDistributed
    var dense2: TimeDistributed
    @noDerivative let dropout: Dropout<Float>

    init(size: Int, hidden: Int, dropProbability: Double) {
        dense1 = TimeDistributed(
            Dense<Float>(inputSize: size, outputSize: hidden, activation: gelu))
        dense2 = TimeDistributed(Dense<Float>(inputSize: hidden, outputSize: size))
        dropout = Dropout<Float>(probability: dropProbability)
    }

    @differentiable(wrt: (self, input))
    func callAsFunction(_ input: Tensor<Float>) -> Tensor<Float> {
        return input.sequenced(through: dense1, dropout, dense2)
    }
}

struct AttentionInput: Differentiable {
    var query: Tensor<Float>
    var key: Tensor<Float>
    var value: Tensor<Float>
}

@differentiable(wrt: (query, key, value), vjp: _vjpMakeAttentionInput)
func makeAttentionInput(query: Tensor<Float>, key: Tensor<Float>, value: Tensor<Float>)
    -> AttentionInput {
    return AttentionInput(query: query, key: key, value: value)
}

func _vjpMakeAttentionInput(query: Tensor<Float>, key: Tensor<Float>, value: Tensor<Float>)
    -> (AttentionInput, (AttentionInput.TangentVector) -> (Tensor<Float>, Tensor<Float>, Tensor<Float>)) {
    let result = AttentionInput(query: query, key: key, value: value)
    return (result, { seed in (seed.query, seed.key, seed.value) })
}

struct AttentionContext: Differentiable {
    var key: Tensor<Float>
    var value: Tensor<Float>
}

@differentiable(wrt: (key, value), vjp: _vjpMakeAttentionContext)
func makeAttentionContext(key: Tensor<Float>, value: Tensor<Float>)
    -> AttentionContext {
    return AttentionContext(key: key, value: value)
}

func _vjpMakeAttentionContext(key: Tensor<Float>, value: Tensor<Float>)
    -> (AttentionContext, (AttentionContext.TangentVector) -> (Tensor<Float>, Tensor<Float>)) {
    let result = AttentionContext(key: key, value: value)
    return (result, { seed in (seed.key, seed.value) })
}

@differentiable(wrt: dotProducts, vjp: _vjpCausallyMasked)
func causallyMasked(_ dotProducts: Tensor<Float>, enable: Bool = false) -> Tensor<Float> {
    if !enable {
        return dotProducts
    }
    let (queryTimeSteps, keyTimeSteps) = (dotProducts.shape[1], dotProducts.shape[2])
    let ones = Tensor<Float>(ones: [1, queryTimeSteps, keyTimeSteps])
    let mask = _Raw.matrixBandPart(
        ones,
        numLower: Tensor(Int32(-1)),
        numUpper: Tensor(Int32(queryTimeSteps - keyTimeSteps)))
    return dotProducts * mask - 1e10 * (1 - mask)
}

// causal mask is intentionally invisible to differentiation
func _vjpCausallyMasked(_ dotProducts: Tensor<Float>, enable: Bool)
    -> (Tensor<Float>, (Tensor<Float>) -> Tensor<Float>) {
    return (causallyMasked(dotProducts, enable: enable), identity)
}

struct Attention: ParameterlessLayer {
    @noDerivative let dropout: Dropout<Float>
    @noDerivative let scale: Tensor<Float>
    @noDerivative let causal: Bool
    
    init(size: Int, causal: Bool = false, dropProbability: Double) {
        scale = Tensor(sqrt(Float(size)))
        dropout = Dropout<Float>(probability: dropProbability)
        self.causal = causal
    }
    
    @differentiable(wrt: (self, input))
    func callAsFunction(_ input: AttentionInput) -> Tensor<Float> {
        var dotProducts = batchedMatmul(input.query, input.key, adjointRight: true)
        dotProducts = causallyMasked(dotProducts, enable: causal) / scale
        return batchedMatmul(dropout(softmax(dotProducts)), input.value)
    }
    
    func callAsFunction(_ input: AttentionInput, state: inout AttentionContext) -> Tensor<Float> {
        state = AttentionContext(
            key: state.key.concatenated(with: input.key, alongAxis: 1),
            value: state.value.concatenated(with: input.value, alongAxis: 1))
        var dotProducts = batchedMatmul(input.query, state.key, adjointRight: true)
        dotProducts = causallyMasked(dotProducts, enable: causal) / scale
        return batchedMatmul(dropout(softmax(dotProducts)), state.value)
    }
}

@differentiable(wrt: input)
func splitHeads(_ input: Tensor<Float>, headCount: Int) -> Tensor<Float> {
    let (batchSize, timeSteps, features) = (input.shape[0], input.shape[1], input.shape[2])
    let featuresPerHead = features / headCount
    let splitLastDim = input.reshaped(to: [batchSize, timeSteps, headCount, featuresPerHead])
    let movedToFront = splitLastDim.transposed(permutation: 0, 2, 1, 3)
    return movedToFront.reshaped(to: [batchSize * headCount, timeSteps, featuresPerHead])
}

@differentiable(wrt: input)
func joinHeads(_ input: Tensor<Float>, headCount: Int) -> Tensor<Float> {
    let (generalizedBatch, timeSteps, featuresPerHead) = (
        input.shape[0], input.shape[1], input.shape[2])
    let batchSize = generalizedBatch / headCount
    let features = featuresPerHead * headCount
    let splitFirstDim = input.reshaped(to: [batchSize, headCount, timeSteps, featuresPerHead])
    let movedToBack = splitFirstDim.transposed(permutation: 0, 2, 1, 3)
    return movedToBack.reshaped(to: [batchSize, timeSteps, features])
}

@differentiable(wrt: input, vjp: _vjpSplitQKV)
func splitQKV(_ input: Tensor<Float>) -> AttentionInput {
    let (generalizedBatch, timeSteps, featuresPerHead) = (
        input.shape[0], input.shape[1], input.shape[2] / 3)
    let query = input.slice(
        lowerBounds: [0, 0, 0],
        upperBounds: [generalizedBatch, timeSteps, featuresPerHead])
    let key = input.slice(
        lowerBounds: [0, 0, featuresPerHead],
        upperBounds: [generalizedBatch, timeSteps, 2 * featuresPerHead])
    let value = input.slice(
        lowerBounds: [0, 0, 2 * featuresPerHead],
        upperBounds: [generalizedBatch, timeSteps, 3 * featuresPerHead])
    return makeAttentionInput(query: query, key: key, value: value)
}

func _vjpSplitQKV(_ input: Tensor<Float>)
    -> (AttentionInput, (AttentionInput.TangentVector) -> Tensor<Float>) {
    let value = splitQKV(input)
    return (value, { seed in
        return _Raw.concatV2([seed.query, seed.key, seed.value], axis: Tensor<Int32>(2))
    })
}

struct MultiHeadAttention: Layer {
    var attention: Attention
    var wqkv: TimeDistributed
    var wo: TimeDistributed
    @noDerivative let headCount: Int
    
    init(attention: Attention, size: Int, headCount: Int) {
        self.attention = attention
        wqkv = TimeDistributed(Dense<Float>(
            inputSize: size, outputSize: size * 3, activation: identity))
        wo = TimeDistributed(Dense<Float>(inputSize: size, outputSize: size, activation: identity))
        self.headCount = headCount
    }
    
    @differentiable(wrt: (self, input))
    func callAsFunction(_ input: Tensor<Float>) -> Tensor<Float> {
        let qkvProjected = wqkv(input)
        let qkvSplit = splitHeads(qkvProjected, headCount: headCount)
        let attentionInput = splitQKV(qkvSplit)
        let outputs = attention(attentionInput)
        return wo(joinHeads(outputs, headCount: headCount))
    }
    
    func callAsFunction(_ input: Tensor<Float>, state: inout AttentionContext) -> Tensor<Float> {
        let qkvProjected = wqkv(input)<commit_msg>Remove _Raw in Transformer/Model.swift (#249)<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.

import TensorFlow
import Python

struct TimeDistributed: Layer {
    var dense: Dense<Float>

    init(_ wrapped: Dense<Float>) {
        self.dense = wrapped
    }

    @differentiable(wrt: (self, input))
    func callAsFunction(_ input: Tensor<Float>) -> Tensor<Float> {
        let (batchSize, timeSteps, features) = (input.shape[0], input.shape[1], input.shape[2])
        let reshaped = input.reshaped(to: [batchSize * timeSteps, features])
        let output = dense(reshaped)
        let outputFeatures = output.shape[1]
        return output.reshaped(to: [batchSize, timeSteps, outputFeatures])
    }
}

struct FeedForward: Layer {
    var dense1: TimeDistributed
    var dense2: TimeDistributed
    @noDerivative let dropout: Dropout<Float>

    init(size: Int, hidden: Int, dropProbability: Double) {
        dense1 = TimeDistributed(
            Dense<Float>(inputSize: size, outputSize: hidden, activation: gelu))
        dense2 = TimeDistributed(Dense<Float>(inputSize: hidden, outputSize: size))
        dropout = Dropout<Float>(probability: dropProbability)
    }

    @differentiable(wrt: (self, input))
    func callAsFunction(_ input: Tensor<Float>) -> Tensor<Float> {
        return input.sequenced(through: dense1, dropout, dense2)
    }
}

struct AttentionInput: Differentiable {
    var query: Tensor<Float>
    var key: Tensor<Float>
    var value: Tensor<Float>
}

@differentiable(wrt: (query, key, value), vjp: _vjpMakeAttentionInput)
func makeAttentionInput(query: Tensor<Float>, key: Tensor<Float>, value: Tensor<Float>)
    -> AttentionInput {
    return AttentionInput(query: query, key: key, value: value)
}

func _vjpMakeAttentionInput(query: Tensor<Float>, key: Tensor<Float>, value: Tensor<Float>)
    -> (AttentionInput, (AttentionInput.TangentVector) -> (Tensor<Float>, Tensor<Float>, Tensor<Float>)) {
    let result = AttentionInput(query: query, key: key, value: value)
    return (result, { seed in (seed.query, seed.key, seed.value) })
}

struct AttentionContext: Differentiable {
    var key: Tensor<Float>
    var value: Tensor<Float>
}

@differentiable(wrt: (key, value), vjp: _vjpMakeAttentionContext)
func makeAttentionContext(key: Tensor<Float>, value: Tensor<Float>)
    -> AttentionContext {
    return AttentionContext(key: key, value: value)
}

func _vjpMakeAttentionContext(key: Tensor<Float>, value: Tensor<Float>)
    -> (AttentionContext, (AttentionContext.TangentVector) -> (Tensor<Float>, Tensor<Float>)) {
    let result = AttentionContext(key: key, value: value)
    return (result, { seed in (seed.key, seed.value) })
}

@differentiable(wrt: dotProducts, vjp: _vjpCausallyMasked)
func causallyMasked(_ dotProducts: Tensor<Float>, enable: Bool = false) -> Tensor<Float> {
    if !enable {
        return dotProducts
    }
    let (queryTimeSteps, keyTimeSteps) = (dotProducts.shape[1], dotProducts.shape[2])
    let ones = Tensor<Float>(ones: [1, queryTimeSteps, keyTimeSteps])
    let mask = ones.bandPart(-1, queryTimeSteps - keyTimeSteps)
    return dotProducts * mask - 1e10 * (1 - mask)
}

// causal mask is intentionally invisible to differentiation
func _vjpCausallyMasked(_ dotProducts: Tensor<Float>, enable: Bool)
    -> (Tensor<Float>, (Tensor<Float>) -> Tensor<Float>) {
    return (causallyMasked(dotProducts, enable: enable), identity)
}

struct Attention: ParameterlessLayer {
    @noDerivative let dropout: Dropout<Float>
    @noDerivative let scale: Tensor<Float>
    @noDerivative let causal: Bool
    
    init(size: Int, causal: Bool = false, dropProbability: Double) {
        scale = Tensor(sqrt(Float(size)))
        dropout = Dropout<Float>(probability: dropProbability)
        self.causal = causal
    }
    
    @differentiable(wrt: (self, input))
    func callAsFunction(_ input: AttentionInput) -> Tensor<Float> {
        var dotProducts = batchedMatmul(input.query, input.key, adjointRight: true)
        dotProducts = causallyMasked(dotProducts, enable: causal) / scale
        return batchedMatmul(dropout(softmax(dotProducts)), input.value)
    }
    
    func callAsFunction(_ input: AttentionInput, state: inout AttentionContext) -> Tensor<Float> {
        state = AttentionContext(
            key: state.key.concatenated(with: input.key, alongAxis: 1),
            value: state.value.concatenated(with: input.value, alongAxis: 1))
        var dotProducts = batchedMatmul(input.query, state.key, adjointRight: true)
        dotProducts = causallyMasked(dotProducts, enable: causal) / scale
        return batchedMatmul(dropout(softmax(dotProducts)), state.value)
    }
}

@differentiable(wrt: input)
func splitHeads(_ input: Tensor<Float>, headCount: Int) -> Tensor<Float> {
    let (batchSize, timeSteps, features) = (input.shape[0], input.shape[1], input.shape[2])
    let featuresPerHead = features / headCount
    let splitLastDim = input.reshaped(to: [batchSize, timeSteps, headCount, featuresPerHead])
    let movedToFront = splitLastDim.transposed(permutation: 0, 2, 1, 3)
    return movedToFront.reshaped(to: [batchSize * headCount, timeSteps, featuresPerHead])
}

@differentiable(wrt: input)
func joinHeads(_ input: Tensor<Float>, headCount: Int) -> Tensor<Float> {
    let (generalizedBatch, timeSteps, featuresPerHead) = (
        input.shape[0], input.shape[1], input.shape[2])
    let batchSize = generalizedBatch / headCount
    let features = featuresPerHead * headCount
    let splitFirstDim = input.reshaped(to: [batchSize, headCount, timeSteps, featuresPerHead])
    let movedToBack = splitFirstDim.transposed(permutation: 0, 2, 1, 3)
    return movedToBack.reshaped(to: [batchSize, timeSteps, features])
}

@differentiable(wrt: input, vjp: _vjpSplitQKV)
func splitQKV(_ input: Tensor<Float>) -> AttentionInput {
    let (generalizedBatch, timeSteps, featuresPerHead) = (
        input.shape[0], input.shape[1], input.shape[2] / 3)
    let query = input.slice(
        lowerBounds: [0, 0, 0],
        upperBounds: [generalizedBatch, timeSteps, featuresPerHead])
    let key = input.slice(
        lowerBounds: [0, 0, featuresPerHead],
        upperBounds: [generalizedBatch, timeSteps, 2 * featuresPerHead])
    let value = input.slice(
        lowerBounds: [0, 0, 2 * featuresPerHead],
        upperBounds: [generalizedBatch, timeSteps, 3 * featuresPerHead])
    return makeAttentionInput(query: query, key: key, value: value)
}

func _vjpSplitQKV(_ input: Tensor<Float>)
    -> (AttentionInput, (AttentionInput.TangentVector) -> Tensor<Float>) {
    let value = splitQKV(input)
    return (value, { seed in
        return Tensor(concatenating: [seed.query, seed.key, seed.value], alongAxis: 2)
    })
}

struct MultiHeadAttention: Layer {
    var attention: Attention
    var wqkv: TimeDistributed
    var wo: TimeDistributed
    @noDerivative let headCount: Int
    
    init(attention: Attention, size: Int, headCount: Int) {
        self.attention = attention
        wqkv = TimeDistributed(Dense<Float>(
            inputSize: size, outputSize: size * 3, activation: identity))
        wo = TimeDistributed(Dense<Float>(inputSize: size, outputSize: size, activation: identity))
        self.headCount = headCount
    }
    
    @differentiable(wrt: (self, input))
    func callAsFunction(_ input: Tensor<Float>) -> Tensor<Float> {
        let qkvProjected = wqkv(input)
        let qkvSplit = splitHeads(qkvProjected, headCount: headCount)
        let attentionInput = splitQKV(qkvSplit)
        let outputs = attention(attentionInput)
        return wo(joinHeads(outputs, headCount: headCount))
    }
    
    func callAsFunction(_ input: Tensor<Float>, state: inout AttentionContext) -> Tensor<Float> {
        let qkvProjected = wqkv(input)
{'fraction_non_alphanumeric': 0.11576369392509409, 'fraction_numerical': 0.007526432112777015, 'mean_word_length': 3.9111176298034613, 'pattern_counts': {'":': 0, '<': 114, '<?xml version=': 0, '>': 154, 'https://': 0, 'lorem ipsum': 0, 'www.': 2, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '4190773', 'n_tokens_mistral': 5267, 'n_tokens_neox': 5062, 'n_words': 1356}