<commit_before>from numpy import arange, dtype, int64, float64
from bolt import array, ones


def test_shape(sc):

    x = arange(2*3).reshape((2, 3))
    b = array(x, sc)
    assert b.shape == x.shape

    x = arange(2*3*4).reshape((2, 3, 4))
    b = array(x, sc)
    assert b.shape == x.shape

def test_size(sc):

    x = arange(2*3*4).reshape((2, 3, 4))
    b = array(x, sc, axis=0)
    assert b.size == x.size

def test_split(sc):

    x = arange(2*3*4).reshape((2, 3, 4))
    b = array(x, sc, axis=0)
    assert b.split == 1

    b = array(x, sc, axis=(0, 1))
    assert b.split == 2

def test_ndim(sc):

    x = arange(2**5).reshape(2, 2, 2, 2, 2)
    b = array(x, sc, axis=(0, 1, 2))

    assert b.keys.ndim == 3
    assert b.values.ndim == 2
    assert b.ndim == 5

def test_mask(sc):

    x = arange(2*3*4).reshape((2, 3, 4))
    b = array(x, sc, axis=0)
    assert b.mask == (1, 0, 0)

    b = array(x, sc, axis=(0, 1))
    assert b.mask == (1, 1, 0)

    b = array(x, sc, axis=(0, 1, 2))
    assert b.mask == (1, 1, 1)

def test_cache(sc):
    x = arange(2*3).reshape((2, 3))
    b = array(x, sc)
    b.cache()
    assert b._rdd.is_cached
    b.unpersist()
    assert not b._rdd.is_cached

def test_dtype(sc):

    a = arange(2**8, dtype=int64)
    b = array(a, sc, dtype=int64)
    assert a.dtype == b.dtype<commit_msg>Add tests for concatenate<commit_after>from numpy import arange, dtype, int64, float64
from bolt import array, ones
from bolt.utils import allclose

def test_shape(sc):

    x = arange(2*3).reshape((2, 3))
    b = array(x, sc)
    assert b.shape == x.shape

    x = arange(2*3*4).reshape((2, 3, 4))
    b = array(x, sc)
    assert b.shape == x.shape

def test_size(sc):

    x = arange(2*3*4).reshape((2, 3, 4))
    b = array(x, sc, axis=0)
    assert b.size == x.size

def test_split(sc):

    x = arange(2*3*4).reshape((2, 3, 4))
    b = array(x, sc, axis=0)
    assert b.split == 1

    b = array(x, sc, axis=(0, 1))
    assert b.split == 2

def test_ndim(sc):

    x = arange(2**5).reshape(2, 2, 2, 2, 2)
    b = array(x, sc, axis=(0, 1, 2))

    assert b.keys.ndim == 3
    assert b.values.ndim == 2
    assert b.ndim == 5

def test_mask(sc):

    x = arange(2*3*4).reshape((2, 3, 4))
    b = array(x, sc, axis=0)
    assert b.mask == (1, 0, 0)

    b = array(x, sc, axis=(0, 1))
    assert b.mask == (1, 1, 0)

    b = array(x, sc, axis=(0, 1, 2))
    assert b.mask == (1, 1, 1)

def test_cache(sc):

    x = arange(2*3).reshape((2, 3))
    b = array(x, sc)
    b.cache()
    assert b._rdd.is_cached
    b.unpersist()
    assert not b._rdd.is_cached

def test_concatenate(sc):

    from numpy import concatenate
    x = arange(2*3).reshape((2, 3))
    b = array(x, sc)
    c = array(x)
    assert allclose(b.concatenate(x).toarray(), concatenate((x, x)))
    assert allclose(b.concatenate(b).toarray(), concatenate((x, x)))
    assert allclose(b.concatenate(c).toarray(), concatenate((x, x)))

def test_dtype(sc):

    a = arange(2**8, dtype=int64)
    b = array(a, sc, dtype=int64)
    assert a.dtype == b.dtype
{'fraction_non_alphanumeric': 0.18661856346343064, 'fraction_numerical': 0.051164316169235814, 'mean_word_length': 2.714981729598051, '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': '4867998', 'n_tokens_mistral': 1394, 'n_tokens_neox': 1298, 'n_words': 338}