<commit_before>#include "xchainer/array.h"

#include <array>
#include <cstddef>
#include <initializer_list>
#include <type_traits>

#ifdef XCHAINER_ENABLE_CUDA
#include <cuda_runtime.h>
#endif  // XCHAINER_ENABLE_CUDA
#include <gtest/gtest.h>

#include "xchainer/array.h"
#ifdef XCHAINER_ENABLE_CUDA
#include "xchainer/cuda/cuda_runtime.h"
#endif  // XCHAINER_ENABLE_CUDA
#include "xchainer/device.h"
#include "xchainer/memory.h"
#include "xchainer/op_node.h"

namespace xchainer {
namespace {

class ArrayTest : public ::testing::TestWithParam<::testing::tuple<std::string>> {
protected:
    void SetUp() override {
        std::string device_name = ::testing::get<0>(GetParam());
        device_scope_ = std::make_unique<DeviceScope>(device_name);
    }

    void TearDown() override { device_scope_.reset(); }

public:
    template <typename T>
    Array MakeArray(std::initializer_list<int64_t> shape, std::shared_ptr<void> data, bool requires_grad = false) {
        Array arr = Array::FromBuffer(shape, TypeToDtype<T>, data);
        arr.set_requires_grad(requires_grad);
        return arr;
    }

    template <typename T>
    Array MakeArray(std::initializer_list<int64_t> shape, std::initializer_list<T> data, bool requires_grad = false) {
        auto a = std::make_unique<T[]>(data.size());
        std::copy(data.begin(), data.end(), a.get());
        return MakeArray<T>(shape, std::move(a), requires_grad);
    }

    template <typename T>
    void ExpectEqual(const Array& expected, const Array& actual) {
        EXPECT_EQ(expected.dtype(), actual.dtype());
        EXPECT_EQ(expected.shape(), actual.shape());
        ExpectDataEqual<T>(expected, actual);
    }

    template <typename T>
    void ExpectDataEqual(const Array& expected, const Array& actual) {
        const T* expected_data = static_cast<const T*>(expected.data().get());
        ExpectDataEqual(expected_data, actual);
    }

    template <typename T>
    void ExpectDataEqual(const T* expected_data, const Array& actual) {
#ifdef XCHAINER_ENABLE_CUDA
        std::string device_name = ::testing::get<0>(GetParam());
        if (device_name == "cuda") {
            cuda::CheckError(cudaDeviceSynchronize());
        }
#endif  // XCHAINER_ENABLE_CUDA
        auto total_size = actual.shape().total_size();
        const T* actual_data = static_cast<const T*>(actual.data().get());
        for (decltype(total_size) i = 0; i < total_size; i++) {
            EXPECT_EQ(expected_data[i], actual_data[i]) << "where i is " << i;
        }
    }

    template <typename T>
    void ExpectDataEqual(T expected, const Array& actual) {
#ifdef XCHAINER_ENABLE_CUDA
        std::string device_name = ::testing::get<0>(GetParam());
        if (device_name == "cuda") {
            cuda::CheckError(cudaDeviceSynchronize());
        }
#endif  // XCHAINER_ENABLE_CUDA
        auto total_size = actual.shape().total_size();
        const T* actual_data = static_cast<const T*>(actual.data().get());
        for (decltype(total_size) i = 0; i < total_size; i++) {
            if (std::isnan(expected)) {
                EXPECT_TRUE(std::isnan(actual_data[i])) << "where i is " << i;
            } else {
                EXPECT_EQ(expected, actual_data[i]) << "where i is " << i;
            }
        }
    }

    void ExpectDataExistsOnCurrentDevice(const Array& array) {
        if (GetCurrentDevice() == MakeDevice("cpu")) {
            EXPECT_FALSE(IsPointerCudaMemory(array.data().get()));
        } else if (GetCurrentDevice() == MakeDevice("cuda")) {
            EXPECT_TRUE(IsPointerCudaMemory(array.data().get()));
        } else {
            FAIL() << "invalid device";
        }
    }

    template <bool is_const, typename T>
    void CheckFromBuffer() {
        using TargetArray = std::conditional_t<is_const, const Array, Array>;

        Shape shape = {3, 2};
        Dtype dtype = TypeToDtype<T>;
        int64_t size = shape.total_size();
        int64_t bytesize = size * sizeof(T);

        auto expected_data = std::make_unique<T[]>(size);
        for (int i = 0; i< size; i++) {
            expected_data[i] = static_cast<T>(i);
        }

        auto data = std::make_unique<T[]>(size);
        std::memcpy(data.get(), expected_data.get(), bytesize);
        auto raw_ptr = data.get();
        TargetArray x = Array::FromBuffer(shape, dtype, std::move(data));

        // Basic attributes
        EXPECT_EQ(shape, x.shape());
        EXPECT_EQ(dtype, x.dtype());
        EXPECT_EQ(2, x.ndim());
        EXPECT_EQ(3 * 2, x.total_size());
        EXPECT_EQ(int64_t{sizeof(T)}, x.element_bytes());
        EXPECT_EQ(bytesize, x.total_bytes());
        EXPECT_EQ(0, x.offset());
        EXPECT_TRUE(x.is_contiguous());

        // Array::data
        ExpectDataEqual<T>(expected_data.get(), x);
        ExpectDataExistsOnCurrentDevice(x);
        if (GetCurrentDevice() == MakeDevice("cpu")) {
            EXPECT_EQ(raw_ptr, x.data().get());
        } else if (GetCurrentDevice() == MakeDevice("cuda")) {
            EXPECT_NE(raw_ptr, x.data().get());
        } else {
            FAIL() << "invalid device";
        }
    }

    template <typename T>
    void CheckEmpty() {
        Dtype dtype = TypeToDtype<T>;
        Array x = Array::Empty(Shape{3, 2}, dtype);
        EXPECT_NE(x.data(), nullptr);
        EXPECT_EQ(x.shape(), Shape({3, 2}));
        EXPECT_EQ(x.dtype(), dtype);
        ExpectDataExistsOnCurrentDevice(x);
    }

    template <typename T>
    void CheckEmptyLike() {
        Dtype dtype = TypeToDtype<T>;
        Array x_orig = Array::Empty(Shape{3, 2}, dtype);
        Array x = Array::EmptyLike(x_orig);
        EXPECT_NE(x.data(), nullptr);
        EXPECT_NE(x.data(), x_orig.data());
        EXPECT_EQ(x.shape(), x_orig.shape());<commit_msg>cosmetic changes<commit_after>#include "xchainer/array.h"

#include <array>
#include <cstddef>
#include <initializer_list>
#include <type_traits>

#ifdef XCHAINER_ENABLE_CUDA
#include <cuda_runtime.h>
#endif  // XCHAINER_ENABLE_CUDA
#include <gtest/gtest.h>

#include "xchainer/array.h"
#ifdef XCHAINER_ENABLE_CUDA
#include "xchainer/cuda/cuda_runtime.h"
#endif  // XCHAINER_ENABLE_CUDA
#include "xchainer/device.h"
#include "xchainer/memory.h"
#include "xchainer/op_node.h"

namespace xchainer {
namespace {

class ArrayTest : public ::testing::TestWithParam<::testing::tuple<std::string>> {
protected:
    void SetUp() override {
        std::string device_name = ::testing::get<0>(GetParam());
        device_scope_ = std::make_unique<DeviceScope>(device_name);
    }

    void TearDown() override { device_scope_.reset(); }

public:
    template <typename T>
    Array MakeArray(std::initializer_list<int64_t> shape, std::shared_ptr<void> data, bool requires_grad = false) {
        Array arr = Array::FromBuffer(shape, TypeToDtype<T>, data);
        arr.set_requires_grad(requires_grad);
        return arr;
    }

    template <typename T>
    Array MakeArray(std::initializer_list<int64_t> shape, std::initializer_list<T> data, bool requires_grad = false) {
        auto a = std::make_unique<T[]>(data.size());
        std::copy(data.begin(), data.end(), a.get());
        return MakeArray<T>(shape, std::move(a), requires_grad);
    }

    template <typename T>
    void ExpectEqual(const Array& expected, const Array& actual) {
        EXPECT_EQ(expected.dtype(), actual.dtype());
        EXPECT_EQ(expected.shape(), actual.shape());
        ExpectDataEqual<T>(expected, actual);
    }

    template <typename T>
    void ExpectDataEqual(const Array& expected, const Array& actual) {
        const T* expected_data = static_cast<const T*>(expected.data().get());
        ExpectDataEqual(expected_data, actual);
    }

    template <typename T>
    void ExpectDataEqual(const T* expected_data, const Array& actual) {
#ifdef XCHAINER_ENABLE_CUDA
        std::string device_name = ::testing::get<0>(GetParam());
        if (device_name == "cuda") {
            cuda::CheckError(cudaDeviceSynchronize());
        }
#endif  // XCHAINER_ENABLE_CUDA
        auto total_size = actual.shape().total_size();
        const T* actual_data = static_cast<const T*>(actual.data().get());
        for (decltype(total_size) i = 0; i < total_size; i++) {
            EXPECT_EQ(expected_data[i], actual_data[i]) << "where i is " << i;
        }
    }

    template <typename T>
    void ExpectDataEqual(T expected, const Array& actual) {
#ifdef XCHAINER_ENABLE_CUDA
        std::string device_name = ::testing::get<0>(GetParam());
        if (device_name == "cuda") {
            cuda::CheckError(cudaDeviceSynchronize());
        }
#endif  // XCHAINER_ENABLE_CUDA
        auto total_size = actual.shape().total_size();
        const T* actual_data = static_cast<const T*>(actual.data().get());
        for (decltype(total_size) i = 0; i < total_size; i++) {
            if (std::isnan(expected)) {
                EXPECT_TRUE(std::isnan(actual_data[i])) << "where i is " << i;
            } else {
                EXPECT_EQ(expected, actual_data[i]) << "where i is " << i;
            }
        }
    }

    void ExpectDataExistsOnCurrentDevice(const Array& array) {
        if (GetCurrentDevice() == MakeDevice("cpu")) {
            EXPECT_FALSE(IsPointerCudaMemory(array.data().get()));
        } else if (GetCurrentDevice() == MakeDevice("cuda")) {
            EXPECT_TRUE(IsPointerCudaMemory(array.data().get()));
        } else {
            FAIL() << "invalid device";
        }
    }

    template <bool is_const, typename T>
    void CheckFromBuffer() {
        using TargetArray = std::conditional_t<is_const, const Array, Array>;

        Shape shape = {3, 2};
        Dtype dtype = TypeToDtype<T>;
        int64_t size = shape.total_size();
        int64_t bytesize = size * sizeof(T);
        T raw_ptr[] = {0, 1, 2, 3, 4, 5};
        auto data = std::shared_ptr<T>(raw_ptr, [](T* ptr) { (void)ptr; });
        TargetArray x = Array::FromBuffer(shape, dtype, data);

        // Basic attributes
        EXPECT_EQ(shape, x.shape());
        EXPECT_EQ(dtype, x.dtype());
        EXPECT_EQ(2, x.ndim());
        EXPECT_EQ(3 * 2, x.total_size());
        EXPECT_EQ(int64_t{sizeof(T)}, x.element_bytes());
        EXPECT_EQ(bytesize, x.total_bytes());
        EXPECT_EQ(0, x.offset());
        EXPECT_TRUE(x.is_contiguous());

        // Array::data
        ExpectDataEqual<T>(raw_ptr, x);
        ExpectDataExistsOnCurrentDevice(x);
        if (GetCurrentDevice() == MakeDevice("cpu")) {
            EXPECT_EQ(raw_ptr, x.data().get());
        } else if (GetCurrentDevice() == MakeDevice("cuda")) {
            EXPECT_NE(raw_ptr, x.data().get());
        } else {
            FAIL() << "invalid device";
        }
    }

    template <typename T>
    void CheckEmpty() {
        Dtype dtype = TypeToDtype<T>;
        Array x = Array::Empty(Shape{3, 2}, dtype);
        EXPECT_NE(x.data(), nullptr);
        EXPECT_EQ(x.shape(), Shape({3, 2}));
        EXPECT_EQ(x.dtype(), dtype);
        ExpectDataExistsOnCurrentDevice(x);
    }

    template <typename T>
    void CheckEmptyLike() {
        Dtype dtype = TypeToDtype<T>;
        Array x_orig = Array::Empty(Shape{3, 2}, dtype);
        Array x = Array::EmptyLike(x_orig);
        EXPECT_NE(x.data(), nullptr);
        EXPECT_NE(x.data(), x_orig.data());
        EXPECT_EQ(x.shape(), x_orig.shape());
{'fraction_non_alphanumeric': 0.14984087694483733, 'fraction_numerical': 0.005392503536067892, 'mean_word_length': 2.9555944055944057, 'pattern_counts': {'":': 0, '<': 118, '<?xml version=': 0, '>': 81, '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': '3182834', 'n_tokens_mistral': 3664, 'n_tokens_neox': 3640, 'n_words': 727}