<commit_before>//
//  checkgradient.swift
//  stem
//
//  Created by <NAME> on 12/8/15.
//  Copyright © 2015 <NAME>. All rights reserved.
//

import Foundation

func calcForwardGrad<S:Storage where S.ElementType:FloatNumericType>
    (op:Op<S>, params:Tensor<S>, eps:Double) -> Tensor<S>
{
    let jacobian:Tensor<S> = zeros(Extent(params.shape.elements, op.output!.shape.elements))

    for i in 0..<params.shape.elements {
        let old_value:S.ElementType = params[i]
        
        // positive direction
        params[i] = old_value + S.ElementType(eps)
        op.apply()
        let pvalue = copy(op.output!)
        
        // negative direction
        params[i] = old_value - S.ElementType(eps)
        op.apply()
        let nvalue = copy(op.output!)
        
        // return to original value
        params[i] = old_value
        
        let diff:Tensor<S> = (pvalue - nvalue) * S.ElementType(1.0/(2.0*eps))
        jacobian[i, all] = ravel(diff)
    }
    
    return jacobian
}

func calcBackwardGrad<S:Storage where S.ElementType:FloatNumericType>
    (forward:Op<S>, _ backward:Op<S>, gradParams:Tensor<S>) -> Tensor<S>
{
    forward.apply()
    backward.apply()
    let gradOutput = ravel(backward.inputs[2].output!)
    
    let jacobian:Tensor<S> = zeros(Extent(gradParams.shape.elements, gradOutput.shape.elements))

    for i in 0..<gradOutput.shape.elements {
        fill(gradOutput, value: 0)
        gradOutput[i] = 1
        
        (backward as! GradientType).reset()
        backward.apply()
        var din = gradParams
        din = din.reshape(Extent(din.shape[0], 1))
        jacobian[all, i] = din
    }
    
    return jacobian
}

public func checkGradient<S:Storage where S.ElementType:FloatNumericType>
    (op:Op<S>, grad:Op<S>, params:Tensor<S>, gradParams:Tensor<S>, eps:Double) -> S.ElementType
{
    op.apply()
    let fgrad = calcForwardGrad(op, params: ravel(params), eps: eps)
    let bgrad = calcBackwardGrad(op, grad, gradParams: ravel(gradParams))

    let error = fgrad - bgrad
    return max(abs(error))
}

public func checkGradient<S:Storage, OpT:Op<S> where S.ElementType:FloatNumericType, OpT:Loss, OpT.StorageType.ElementType==S.ElementType>
    (op:OpT, grad:Op<S>, input:Tensor<S>, eps:Double) -> S.ElementType
{
    op.apply()
    grad.apply()
    
    //    let result = Tensor<S>(Extent(gradParams.map { $0.shape.elements }.reduce(0, combine: +)))
    let result:Tensor<S> = zeros(grad.output!.shape)

    // copy gradients, they will be overwritten from subsequent calls to `fn`
    let analytical_diff = copy(grad.output!)
    
    for i in 0..<input.shape.elements {
        let old_value:S.ElementType = input[i]
        
        // positive direction
        input[i] = old_value + S.ElementType(eps)
        op.apply()
        let pvalue = op.value
        
        // negative direction
        input[i] = old_value - S.ElementType(eps)
        op.apply()
        let nvalue = op.value
        
        
        // return to original value
        input[i] = old_value
        
        let numerical_diff = (pvalue - nvalue) / S.ElementType(2.0*eps)
        
        // TODO: Look into alternate formulations (e.g. either norm of both, or max of denom.)
        result[i] = abs((numerical_diff - analytical_diff[i])/analytical_diff[i])
    }
    
    return max(abs(result))
}


//public func checkGradient<S:Storage where S.ElementType:FloatNumericType>
//    (params params:Tensor<S>,
//            gradParams:Tensor<S>,
//            eps:Double,
//            fn:() -> S.ElementType) -> Tensor<S>
//{
////    let result = Tensor<S>(Extent(gradParams.map { $0.shape.elements }.reduce(0, combine: +)))
//    let result:Tensor<S> = zeros(Extent(gradParams.shape))
//    
//    // calculate gradients for `analytical_diff`
//    fn()
//    
//    // copy gradients, they will be overwritten from subsequent calls to `fn`
//    let analytical_diff = copy(gradParams)
//    
//    for i in 0..<params.shape[0] {
//        let old_value:S.ElementType = params[i]
//        
//        // positive direction
//        params[i] = old_value + S.ElementType(eps)
//        let pvalue = fn()
//        
//        // negative direction
//        params[i] = old_value - S.ElementType(eps)
//        let nvalue = fn()
//        
//        // return to original value
//        params[i] = old_value
//        
//        let numerical_diff = (pvalue - nvalue) / S.ElementType(2.0*eps)
//        
//        // TODO: Look into alternate formulations (e.g. either norm of both, or max of denom.)
//        result[i] = abs((numerical_diff - analytical_diff[i])/analytical_diff[i])
//    }
//    
//    return result
//}
//
//public func checkGradient<S:Storage where S.ElementType:FloatNumericType>
//    (params params:[Tensor<S>],
//     gradParams:[Tensor<S>],
//     eps:Double,
//     fn:() -> S.ElementType) -> Tensor<S>
//{
//    let result = Tensor<S>(Extent(gradParams.map { $0.shape.elements }.reduce(0, combine: +)))
//    
//    // calculate gradients for `analytical_diff`
//    fn()
//    
//    // copy gradients, they will be overwritten from subsequent calls to `fn`
////    let analytical_diff = copy(gradParams)
//    print(gradParams.map { copy($0) })
//    let analytical_diff = concat(gradParams.map { copy($0).reshape(Extent($0.shape.elements)) }, axis: 0)
//    print("adiff = \(analytical_diff)")
//    
//    var c = 0
//    
//    for param in params {
//        let p = param.reshape(Extent(param.shape.elements))
////        print("p = \(p)")
//        for i in 0..<p.shape[0] {
//            let old_value:S.ElementType = p[i]
//            
//            // positive direction
//            p[i] = old_value + S.ElementType(eps)
//            let pvalue = fn()
//            
//            // negative direction
//            p[i] = old_value - S.ElementType(eps)
//            let nvalue = fn()
//            
//            // return to original value
//            p[i] = old_value
//            
//            let numerical_diff = (pvalue - nvalue) / S.ElementType(2.0*eps)
//            
//            // TODO: Look into alternate formulations (e.g. either norm of both, or max of denom.)
//            result[c] = abs((numerical_diff - analytical_diff[c])/analytical_diff[c])
//            c += 1
//        }
//    }
//    
//    return result
//}<commit_msg>Old checkgradient code removed.<commit_after>//
//  checkgradient.swift
//  stem
//
//  Created by <NAME> on 12/8/15.
//  Copyright © 2015 <NAME>. All rights reserved.
//

import Foundation

func calcForwardGrad<S:Storage where S.ElementType:FloatNumericType>
    (op:Op<S>, params:Tensor<S>, eps:Double) -> Tensor<S>
{
    let jacobian:Tensor<S> = zeros(Extent(params.shape.elements, op.output!.shape.elements))

    for i in 0..<params.shape.elements {
        let old_value:S.ElementType = params[i]
        
        // positive direction
        params[i] = old_value + S.ElementType(eps)
        op.apply()
        let pvalue = copy(op.output!)
        
        // negative direction
        params[i] = old_value - S.ElementType(eps)
        op.apply()
        let nvalue = copy(op.output!)
        
        // return to original value
        params[i] = old_value
        
        let diff:Tensor<S> = (pvalue - nvalue) * S.ElementType(1.0/(2.0*eps))
        jacobian[i, all] = ravel(diff)
    }
    
    return jacobian
}

func calcBackwardGrad<S:Storage where S.ElementType:FloatNumericType>
    (forward:Op<S>, _ backward:Op<S>, gradParams:Tensor<S>) -> Tensor<S>
{
    forward.apply()
    backward.apply()
    let gradOutput = ravel(backward.inputs[2].output!)
    
    let jacobian:Tensor<S> = zeros(Extent(gradParams.shape.elements, gradOutput.shape.elements))

    for i in 0..<gradOutput.shape.elements {
        fill(gradOutput, value: 0)
        gradOutput[i] = 1
        
        (backward as! GradientType).reset()
        backward.apply()
        var din = gradParams
        din = din.reshape(Extent(din.shape[0], 1))
        jacobian[all, i] = din
    }
    
    return jacobian
}

public func checkGradient<S:Storage where S.ElementType:FloatNumericType>
    (op:Op<S>, grad:Op<S>, params:Tensor<S>, gradParams:Tensor<S>, eps:Double) -> S.ElementType
{
    op.apply()
    let fgrad = calcForwardGrad(op, params: ravel(params), eps: eps)
    let bgrad = calcBackwardGrad(op, grad, gradParams: ravel(gradParams))

    let error = fgrad - bgrad
    return max(abs(error))
}

public func checkGradient<S:Storage, OpT:Op<S> where S.ElementType:FloatNumericType, OpT:Loss, OpT.StorageType.ElementType==S.ElementType>
    (op:OpT, grad:Op<S>, input:Tensor<S>, eps:Double) -> S.ElementType
{
    op.apply()
    grad.apply()
    
    //    let result = Tensor<S>(Extent(gradParams.map { $0.shape.elements }.reduce(0, combine: +)))
    let result:Tensor<S> = zeros(grad.output!.shape)

    // copy gradients, they will be overwritten from subsequent calls to `fn`
    let analytical_diff = copy(grad.output!)
    
    for i in 0..<input.shape.elements {
        let old_value:S.ElementType = input[i]
        
        // positive direction
        input[i] = old_value + S.ElementType(eps)
        op.apply()
        let pvalue = op.value
        
        // negative direction
        input[i] = old_value - S.ElementType(eps)
        op.apply()
        let nvalue = op.value
        
        
        // return to original value
        input[i] = old_value
        
        let numerical_diff = (pvalue - nvalue) / S.ElementType(2.0*eps)
        
        // TODO: Look into alternate formulations (e.g. either norm of both, or max of denom.)
        result[i] = abs((numerical_diff - analytical_diff[i])/analytical_diff[i])
    }
    
    return max(abs(result))
}
{'fraction_non_alphanumeric': 0.14442732408834102, 'fraction_numerical': 0.006985105290190036, 'mean_word_length': 2.1910848901999342, 'pattern_counts': {'":': 0, '<': 72, '<?xml version=': 0, '>': 76, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, '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': '6167588', 'n_tokens_mistral': 3257, 'n_tokens_neox': 3163, 'n_words': 680}