<commit_before># solve.py

# Imports
import matplotlib.pyplot as plt
import numpy as np

from scipy.optimize import fmin_cg
from shard import Shard

# subaxes
def subaxes(n):
    rows = int(np.floor(np.sqrt(n)))
    cols = int(np.ceil(float(n) / rows))
    w = 1.0 / cols
    h = 1.0 / rows

    f = plt.figure()
    axs = []
    for i in xrange(n):
        r = (rows - 1) - i / cols
        c = i % cols
        ax = f.add_axes([c * w, r * h, w, h], frameon=False)
        axs.append(ax)

    if len(axs) == 1:
        axs = axs[0]

    return f, axs

# shard_gradient
def shard_gradient(I, J, alpha, X, y, k, epsilon=1e-6):
    shape = I.shape[:2]
    domain = shape[::-1]
    shard = Shard(X, k)

    H, dX = shard(domain, return_dX=True, epsilon=epsilon)

    d = alpha * (J - y)
    R = (I - J) + d * H[..., np.newaxis]
    J = dX[..., np.newaxis] * d
    J = J.reshape(X.size, -1).transpose()
    dX = np.dot(R.ravel(), J)
    return dX.reshape(X.shape)

# fit_shard
def fit_shard(I, J, alpha, X, y, k, epsilon=1e-6, update_colours=False,
              limit_colours=True, **kwargs):
    shape = I.shape[:2]
    domain = shape[::-1]

    R0 = (I - J)
    y = np.copy(y)

    def f(x):
        shard = Shard(x.reshape(X.shape), k)
        H = shard(domain)
        R = R0 + alpha * (J - y) * H[..., np.newaxis]
        r = R.ravel()
        return np.dot(r, r)

    def fprime(x, return_energy=False):
        shard = Shard(x.reshape(X.shape), k)
        H, dX = shard(domain, return_dX=True, epsilon=epsilon)
        d = alpha * (J - y)
        R = R0 + d * H[..., np.newaxis]
        J_ = dX[..., np.newaxis] * d
        J_ = J_.reshape(X.size, -1).transpose()
        r = R.ravel()
        dx = np.dot(r, J_)

        if not return_energy:
            return dx
        else:
            return np.dot(r, r), dx

    callbacks = []
    def callback_handler(xk):
        for callback in callbacks:<commit_msg>Add `check_gradients` argument to `fit_shard` in solve.py<commit_after># solve.py

# Imports
import matplotlib.pyplot as plt
import numpy as np

from misc.scipy_ import approx_jac, print_comparison
from scipy.optimize import fmin_cg
from shard import Shard

# subaxes
def subaxes(n):
    rows = int(np.floor(np.sqrt(n)))
    cols = int(np.ceil(float(n) / rows))
    w = 1.0 / cols
    h = 1.0 / rows

    f = plt.figure()
    axs = []
    for i in xrange(n):
        r = (rows - 1) - i / cols
        c = i % cols
        ax = f.add_axes([c * w, r * h, w, h], frameon=False)
        axs.append(ax)

    if len(axs) == 1:
        axs = axs[0]

    return f, axs

# shard_gradient
def shard_gradient(I, J, alpha, X, y, k, epsilon=1e-6):
    shape = I.shape[:2]
    domain = shape[::-1]
    shard = Shard(X, k)

    H, dX = shard(domain, return_dX=True, epsilon=epsilon)

    d = alpha * (J - y)
    R = (I - J) + d * H[..., np.newaxis]
    J = dX[..., np.newaxis] * d
    J = J.reshape(X.size, -1).transpose()
    dX = np.dot(R.ravel(), J)
    return dX.reshape(X.shape)

# fit_shard
def fit_shard(I, J, alpha, X, y, k, epsilon=1e-6, update_colours=False,
              limit_colours=True, check_gradients=False, **kwargs):
    shape = I.shape[:2]
    domain = shape[::-1]

    R0 = (I - J)
    y = np.copy(y)

    def f(x):
        shard = Shard(x.reshape(X.shape), k)
        H = shard(domain)
        R = R0 + alpha * (J - y) * H[..., np.newaxis]
        r = R.ravel()
        return np.dot(r, r)

    def fprime(x, return_energy=False):
        shard = Shard(x.reshape(X.shape), k)
        H, dX = shard(domain, return_dX=True, epsilon=epsilon)
        d = alpha * (J - y)
        R = R0 + d * H[..., np.newaxis]
        J_ = dX[..., np.newaxis] * d
        J_ = J_.reshape(X.size, -1).transpose()
        r = R.ravel()
        dx = np.dot(r, J_)

        if not return_energy:
            return dx
        else:
            return np.dot(r, r), dx

    if check_gradients:
        x = X.ravel()
        approx_D = approx_jac(lambda x: np.r_[f(x)], x, epsilon=1e-6)[0]
        D = fprime(x)
        print_comparison(approx_D=approx_D, D=D, atol=1e-3)

    callbacks = []
    def callback_handler(xk):
        for callback in callbacks:
{'fraction_non_alphanumeric': 0.16129032258064516, 'fraction_numerical': 0.010832932113625422, 'mean_word_length': 1.9678571428571427, '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': '2571421', 'n_tokens_mistral': 1726, 'n_tokens_neox': 1677, 'n_words': 420}