<commit_before>from __future__ import absolute_import
import scipy.signal

from autograd.core import getval, primitive
from autograd.numpy import flipud

convolve = primitive(scipy.signal.convolve)
convolve2d = primitive(scipy.signal.convolve2d)

def get_same_slice(L_in0, L_in1):
    left_pad = L_in0 - (L_in1 + 1) / 2
    return slice(left_pad, left_pad + L_in1)

def make_grad_convolve_0(ans, in0, in1, mode='full'):
    if mode == 'full':
        return lambda g: convolve(g, flipud(in1), mode='valid')
    elif mode == 'same':
        return lambda g: flipud(convolve(flipud(g), in1, mode='same'))
    elif mode == 'valid':
        return lambda g: convolve(g, flipud(in1), mode='full')
    else:
        raise Exception("Unrecognized mode {0}".format(mode))

convolve.defgrad(make_grad_convolve_0, argnum=0)

def make_grad_convolve_1(ans, in0, in1, mode='full'):
    if mode == 'full':
        return lambda g: flipud(convolve(flipud(g), in0, mode='valid'))
    elif mode == 'same':
        def grad_fun(g):
            idxs = get_same_slice(in0.shape[0], in1.shape[0])
            return  flipud(convolve(in0, flipud(g), mode='full'))[idxs]
        return grad_fun
    elif mode == 'valid':
        return lambda g: flipud(convolve(in0, flipud(g), mode='valid'))
    else:
        raise Exception("Unrecognized mode {0}".format(mode))

convolve.defgrad(make_grad_convolve_1, argnum=1)

def flip2d(x):
    return x[::-1, ::-1]

def make_grad_convolve2d_0(ans, in0, in1, mode='full'):
    if mode == 'full':
        return lambda g: convolve2d(g, flip2d(in1), mode='valid')
    elif mode == 'same':
        return lambda g: flip2d(convolve2d(flip2d(g), in1, mode='same'))
    elif mode == 'valid':
        return lambda g: convolve2d(g, flip2d(in1), mode='full')
    else:
        raise Exception("Unrecognized mode {0}".format(mode))

convolve2d.defgrad(make_grad_convolve2d_0, argnum=0)

def make_grad_convolve2d_1(ans, in0, in1, mode='full'):
    if mode == 'full':
        return lambda g: flip2d(convolve2d(flip2d(g), in0, mode='valid'))
    elif mode == 'same':
        def grad_fun(g):
            idxs_0 = get_same_slice(in0.shape[0], in1.shape[0])
            idxs_1 = get_same_slice(in0.shape[1], in1.shape[1])
            return  flip2d(convolve(in0, flip2d(g), mode='full'))[idxs_0, idxs_1]
        return grad_fun
    elif mode == 'valid':
        return lambda g: flip2d(convolve2d(in0, flip2d(g), mode='valid'))
    else:
        raise Exception("Unrecognized mode {0}".format(mode))

convolve2d.defgrad(make_grad_convolve2d_1, argnum=1)<commit_msg>Fixed mode='same' grad for convolve2d and reduced the number of flips needed<commit_after>from __future__ import absolute_import
import scipy.signal

from autograd.core import getval, primitive
from autograd.numpy import flipud

convolve = primitive(scipy.signal.convolve)
convolve2d = primitive(scipy.signal.convolve2d)

def get_same_slice(L_in0, L_in1):
    left_pad = L_in0 - (L_in1 + 1) / 2
    return slice(left_pad, left_pad + L_in1)

def make_grad_convolve_0(ans, in0, in1, mode='full'):
    if mode == 'full':
        return lambda g: convolve(g, flipud(in1), mode='valid')
    elif mode == 'same':
        return lambda g: flipud(convolve(flipud(g), in1, mode='same'))
    elif mode == 'valid':
        return lambda g: convolve(g, flipud(in1), mode='full')
    else:
        raise Exception("Unrecognized mode {0}".format(mode))

convolve.defgrad(make_grad_convolve_0, argnum=0)

def make_grad_convolve_1(ans, in0, in1, mode='full'):
    if mode == 'full':
        return lambda g: convolve(g, flipud(in0), mode='valid')
    elif mode == 'same':
        idxs = get_same_slice(in0.shape[0], in1.shape[0])
        return lambda g : convolve(g, flipud(in0), mode='full')[idxs]
    elif mode == 'valid':
        return lambda g: convolve(flipud(in0), g, mode='valid')
    else:
        raise Exception("Unrecognized mode {0}".format(mode))

convolve.defgrad(make_grad_convolve_1, argnum=1)

def flip2d(x):
    return x[::-1, ::-1]

def make_grad_convolve2d_0(ans, in0, in1, mode='full'):
    if mode == 'full':
        return lambda g: convolve2d(g, flip2d(in1), mode='valid')
    elif mode == 'same':
        return lambda g: flip2d(convolve2d(flip2d(g), in1, mode='same'))
    elif mode == 'valid':
        return lambda g: convolve2d(g, flip2d(in1), mode='full')
    else:
        raise Exception("Unrecognized mode {0}".format(mode))

convolve2d.defgrad(make_grad_convolve2d_0, argnum=0)

def make_grad_convolve2d_1(ans, in0, in1, mode='full'):
    if mode == 'full':
        return lambda g: convolve2d(g, flip2d(in0), mode='valid')
    elif mode == 'same':
        idxs_0 = get_same_slice(in0.shape[0], in1.shape[0])
        idxs_1 = get_same_slice(in0.shape[1], in1.shape[1])
        return lambda g : convolve2d(g, flip2d(in0), mode='full')[idxs_0, idxs_1]
    elif mode == 'valid':
        return lambda g: flip2d(convolve2d(in0, flip2d(g), mode='valid'))
    else:
        raise Exception("Unrecognized mode {0}".format(mode))

convolve2d.defgrad(make_grad_convolve2d_1, argnum=1)
{'fraction_non_alphanumeric': 0.14550839091806514, 'fraction_numerical': 0.03356367226061204, 'mean_word_length': 3.8339694656488548, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, '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': '6403419', 'n_tokens_mistral': 1954, 'n_tokens_neox': 1876, 'n_words': 418}