<commit_before>
#[macro_use(s)]
extern crate ndarray;
extern crate ndarray_rand;
extern crate rand;

use ndarray_rand::{RandomExt, F32};
use rand::Rng;

use ndarray::{
    ArrayBase,
    Array,
    Dimension,
    Data,
    DataMut,
    Axis, RemoveAxis,
    LinalgScalar,
    Ix,
};

use rand::distributions::Normal;

// simple, slow, correct (hopefully) mat mul
fn reference_mat_mul<A, S, S2>(lhs: &ArrayBase<S, (Ix, Ix)>, rhs: &ArrayBase<S2, (Ix, Ix)>)
    -> Array<A, (Ix, Ix)>
    where A: LinalgScalar,
          S: Data<Elem=A>,
          S2: Data<Elem=A>,
{
    let ((m, k), (_, n)) = (lhs.dim(), rhs.dim());
    let mut res_elems = Vec::<A>::with_capacity(m * n);
    unsafe {
        res_elems.set_len(m * n);
    }

    let mut i = 0;
    let mut j = 0;
    for rr in &mut res_elems {
        unsafe {
            *rr = (0..k).fold(A::zero(),
                move |s, x| s + *lhs.uget((i, x)) * *rhs.uget((x, j)));
        }
        j += 1;
        if j == n {
            j = 0;
            i += 1;
        }
    }
    unsafe {
        ArrayBase::from_vec_dim_unchecked((m, n), res_elems)
    }
}

fn gen<D>(d: D) -> Array<f32, D>
    where D: Dimension,
{
    Array::random(d, F32(Normal::new(0., 1.)))
}
fn gen_f64<D>(d: D) -> Array<f64, D>
    where D: Dimension,
{
    Array::random(d, Normal::new(0., 1.))
}

#[test]
fn accurate_eye_f32() {
    for i in 0..20 {
        let eye = Array::eye(i);
        for j in 0..20 {
            let a = gen((i, j));
            let a2 = eye.dot(&a);
            if !a.all_close(&a2, 1e-6) {
                panic!("Arrays are not equal:\n{:?}\n{:?}\n{:?}", a, a2, &a2 - &a);
            }
            let a3 = a.t().dot(&eye);
            if !a.t().all_close(&a3, 1e-6) {
                panic!("Arrays are not equal:\n{:?}\n{:?}\n{:?}", a.t(), a3, &a3 - &a.t());
            }
        }<commit_msg>Update ndarray-tests for ndarray 0.6<commit_after>
#[macro_use(s)]
extern crate ndarray;
extern crate ndarray_rand;
extern crate rand;

use ndarray_rand::{RandomExt, F32};
use rand::Rng;

use ndarray::{
    ArrayBase,
    Array,
    Dimension,
    Data,
    DataMut,
    Axis, RemoveAxis,
    LinalgScalar,
    Ix,
};

use rand::distributions::Normal;

// simple, slow, correct (hopefully) mat mul
fn reference_mat_mul<A, S, S2>(lhs: &ArrayBase<S, (Ix, Ix)>, rhs: &ArrayBase<S2, (Ix, Ix)>)
    -> Array<A, (Ix, Ix)>
    where A: LinalgScalar,
          S: Data<Elem=A>,
          S2: Data<Elem=A>,
{
    let ((m, k), (_, n)) = (lhs.dim(), rhs.dim());
    let mut res_elems = Vec::<A>::with_capacity(m * n);
    unsafe {
        res_elems.set_len(m * n);
    }

    let mut i = 0;
    let mut j = 0;
    for rr in &mut res_elems {
        unsafe {
            *rr = (0..k).fold(A::zero(),
                move |s, x| s + *lhs.uget((i, x)) * *rhs.uget((x, j)));
        }
        j += 1;
        if j == n {
            j = 0;
            i += 1;
        }
    }
    unsafe {
        ArrayBase::from_shape_vec_unchecked((m, n), res_elems)
    }
}

fn gen<D>(d: D) -> Array<f32, D>
    where D: Dimension,
{
    Array::random(d, F32(Normal::new(0., 1.)))
}
fn gen_f64<D>(d: D) -> Array<f64, D>
    where D: Dimension,
{
    Array::random(d, Normal::new(0., 1.))
}

#[test]
fn accurate_eye_f32() {
    for i in 0..20 {
        let eye = Array::eye(i);
        for j in 0..20 {
            let a = gen((i, j));
            let a2 = eye.dot(&a);
            if !a.all_close(&a2, 1e-6) {
                panic!("Arrays are not equal:\n{:?}\n{:?}\n{:?}", a, a2, &a2 - &a);
            }
            let a3 = a.t().dot(&eye);
            if !a.t().all_close(&a3, 1e-6) {
                panic!("Arrays are not equal:\n{:?}\n{:?}\n{:?}", a.t(), a3, &a3 - &a.t());
            }
        }
{'fraction_non_alphanumeric': 0.1995172968624296, 'fraction_numerical': 0.02359882005899705, 'mean_word_length': 1.9117876658860264, 'pattern_counts': {'":': 0, '<': 25, '<?xml version=': 0, '>': 31, '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': '3722272', 'n_tokens_mistral': 1568, 'n_tokens_neox': 1521, 'n_words': 354}