<commit_before>import argparse
import contextlib
import time

import numpy as np
from scipy.sparse.linalg import cg
import six

import cupy


@contextlib.contextmanager
def timer(message):
    cupy.cuda.Stream.null.synchronize()
    start = time.time()
    yield
    cupy.cuda.Stream.null.synchronize()
    end = time.time()
    print('%s:  %f sec' % (message, end - start))


def fit(A, b, x0, tol, max_iter):
    xp = cupy.get_array_module(A)
    x = x0
    r0 = b - xp.dot(A, x)
    p = r0
    for i in six.moves.range(max_iter):
        a = xp.dot(r0.T, r0) / xp.dot(xp.dot(p.T, A), p)
        x = x + p * a
        r1 = r0 - xp.dot(A * a, p)
        if xp.linalg.norm(r1) < tol:
            return x
        b = xp.dot(r1.T, r1) / xp.dot(r0.T, r0)
        p = r1 + b * p
        r0 = r1
    msg = 'Failed to converge. Increase max-iter or tol.'
    print(msg)
    return x


def run(gpuid, tol, max_iter):
    # Solve simultaneous linear equations, Ax = b.
    # To check whether the answers are correct, each 'b' is printed.
    for repeat in range(3):
        print("Trial: %d" % repeat)
        # Create the large sparse symmetric matrix 'A'.
        N = 2000
        A = np.random.randint(2, size=(N, N))
        A = ((A + A.T) % 2).astype(np.float32)
        b = A[:, N - 1]
        x0 = np.zeros(N, dtype=np.float32)

        np.set_printoptions(precision=2, suppress=True)
        print('b[:18]=')
        print(b[:18])

        with timer(' CPU '):
            print('Running CPU...')
            x = fit(A, b, x0, tol, max_iter)
            b_calc = np.dot(A, x)
            print(b_calc[:18])

        with cupy.cuda.Device(gpuid):
            print('Running GPU...')
            A_gpu = cupy.asarray(A)
            b_gpu = cupy.asarray(b)
            x0_gpu = cupy.asarray(x0)
            with timer(' GPU '):
                x = fit(A_gpu, b_gpu, x0_gpu, tol, max_iter)
                b_calc = cupy.dot(A_gpu, x)
                print(b_calc[:18])

        with timer(' SciPy '):
            print('Running SciPy...')
            x = cg(A, b)
            b_calc = np.dot(A, x[0])
            print(b_calc[:18])

        print()


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--gpu-id', '-g', default=0, type=int,
                        help='ID of GPU.')
    parser.add_argument('--tol', '-t', default=0.1, type=float,
                        help='tolerance to stop iteration')
    parser.add_argument('--max-iter', '-m', default=5000, type=int,
                        help='number of iterations')
    args = parser.parse_args()
    run(args.gpu_id, args.tol, args.max_iter)<commit_msg>fix doc and creation<commit_after>import argparse
import contextlib
import time

import numpy as np
from scipy.sparse.linalg import cg
import six

import cupy


@contextlib.contextmanager
def timer(message):
    cupy.cuda.Stream.null.synchronize()
    start = time.time()
    yield
    cupy.cuda.Stream.null.synchronize()
    end = time.time()
    print('%s:  %f sec' % (message, end - start))


def fit(A, b, x0, tol, max_iter):
    xp = cupy.get_array_module(A)
    x = x0
    r0 = b - xp.dot(A, x)
    p = r0
    for i in six.moves.range(max_iter):
        a = xp.dot(r0.T, r0) / xp.dot(xp.dot(p.T, A), p)
        x = x + p * a
        r1 = r0 - xp.dot(A * a, p)
        if xp.linalg.norm(r1) < tol:
            return x
        b = xp.dot(r1.T, r1) / xp.dot(r0.T, r0)
        p = r1 + b * p
        r0 = r1
    msg = 'Failed to converge. Increase max-iter or tol.'
    print(msg)
    return x


def run(gpu_id, tol, max_iter):
    '''CuPy Congugate gradient example

    Solve simultaneous linear equations, Ax = b.
    'x' is computed in three ways. To check whether the answers are correct,
    'b' is computed from matrix multiplication of 'A' and 'x' in each case,
    and printed.

    '''
    for repeat in range(3):
        print("Trial: %d" % repeat)
        # Create the large symmetric matrix 'A'.
        N = 2000
        A = np.random.randint(50, size=(N, N))
        A = (A + A.T).astype(np.float64)
        b = np.random.randint(50, size=N).astype(np.float64)
        x0 = np.zeros(N, dtype=np.float64)

        np.set_printoptions(precision=2, suppress=True)
        print('b[:18]=')
        print(b[:18])

        print('Running CPU...')
        with timer(' CPU '):
            x = fit(A, b, x0, tol, max_iter)
        b_calc = np.dot(A, x)
        print(b_calc[:18])

        with cupy.cuda.Device(gpu_id):
            A_gpu = cupy.asarray(A)
            b_gpu = cupy.asarray(b)
            x0_gpu = cupy.asarray(x0)
            print('Running GPU...')
            with timer(' GPU '):
                x = fit(A_gpu, b_gpu, x0_gpu, tol, max_iter)
            b_calc = cupy.dot(A_gpu, x)
            print(b_calc[:18])

        print('Running SciPy...')
        with timer(' SciPy '):
            x = cg(A, b)
        b_calc = np.dot(A, x[0])
        print(b_calc[:18])

        print()


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--gpu-id', '-g', default=0, type=int,
                        help='ID of GPU.')
    parser.add_argument('--tol', '-t', default=0.1, type=float,
                        help='tolerance to stop iteration')
    parser.add_argument('--max-iter', '-m', default=5000, type=int,
                        help='number of iterations')
    args = parser.parse_args()
    run(args.gpu_id, args.tol, args.max_iter)
{'fraction_non_alphanumeric': 0.14589442815249268, 'fraction_numerical': 0.019611436950146627, 'mean_word_length': 1.8949602122015916, 'pattern_counts': {'":': 0, '<': 5, '<?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': '2800730', 'n_tokens_mistral': 2051, 'n_tokens_neox': 1966, 'n_words': 490}