<commit_before>#include "xchainer/routines/manipulation.h"

#include <algorithm>
#include <cstdint>
#include <memory>
#include <string>
#include <vector>

#include <gtest/gtest.h>
#include <nonstd/optional.hpp>

#include "xchainer/array.h"
#include "xchainer/array_index.h"
#include "xchainer/check_backward.h"
#include "xchainer/device_id.h"
#include "xchainer/error.h"
#include "xchainer/shape.h"
#include "xchainer/strides.h"
#include "xchainer/testing/array.h"
#include "xchainer/testing/array_check.h"
#include "xchainer/testing/device_session.h"

namespace xchainer {
namespace {

class ManipulationTest : public ::testing::TestWithParam<std::string> {
protected:
    void SetUp() override {
        const std::string& backend_name = GetParam();
        device_session_.emplace(DeviceId{backend_name, 0});
    }

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

private:
    nonstd::optional<testing::DeviceSession> device_session_;
};

TEST_P(ManipulationTest, AsScalar) {
    using T = float;
    T value = 2.0f;
    Array a = testing::BuildArray<T>({1, 1, 1}, {value}).WithPadding(1);
    Scalar s = AsScalar(a);

    ASSERT_EQ(s.dtype(), TypeToDtype<T>);
    EXPECT_EQ(static_cast<T>(s), value);
}

TEST_P(ManipulationTest, AsScalarInvalidZeroElement) {
    Array a = testing::BuildArray<float>({0}, {});
    EXPECT_THROW(AsScalar(a), DimensionError);
}

TEST_P(ManipulationTest, AsScalarInvalidMoreThanOneElements) {
    Array a = testing::BuildArray<float>({2}, {1.0f, 2.0f});
    EXPECT_THROW(AsScalar(a), DimensionError);
}

TEST_P(ManipulationTest, Transpose) {
    Array a = testing::BuildArray({2, 3})         //
                      .WithLinearData<int32_t>()  //
                      .WithPadding(0);
    Array b = Transpose(a);

    EXPECT_EQ(Shape({3, 2}), b.shape());
    EXPECT_EQ(Strides({4, 12}), b.strides());

    Array e = testing::BuildArray({3, 2}).WithData<int32_t>({0, 3, 1, 4, 2, 5});
    testing::ExpectEqual(e, b);
}

TEST_P(ManipulationTest, TransposeNoncontiguous) {
    Array a = testing::BuildArray({2, 3})         //
                      .WithLinearData<int32_t>()  //
                      .WithPadding(1);
    Array b = Transpose(a);

    EXPECT_EQ(Shape({3, 2}), b.shape());

    Array e = testing::BuildArray({3, 2}).WithData<int32_t>({0, 3, 1, 4, 2, 5});
    testing::ExpectEqual(e, b);
}

TEST_P(ManipulationTest, TransposeBackward) {
    CheckBackwardComputation(
            [](const std::vector<Array>& xs) -> std::vector<Array> { return {Transpose(xs[0])}; },
            {Array::Zeros({2, 3}, Dtype::kFloat32).RequireGrad()},
            {testing::BuildArray({3, 2}, {1.f, 2.f, 3.f, 4.f, 5.f, 6.f})},
            {Array::Full({2, 3}, 1e-5f)});
}

TEST_P(ManipulationTest, TransposeDoubleBackward) {
    CheckDoubleBackwardComputation(
            [](const std::vector<Array>& xs) -> std::vector<Array> {
                auto t = Transpose(xs[0]);
                return {t * t};  // to make it nonlinear
            },
            {(*testing::BuildArray({2, 3}, {1.f, -1.f, 2.f, -2.f, 3.f, -3.f})).RequireGrad()},
            {Array::Ones({3, 2}, Dtype::kFloat32).RequireGrad()},
            {Array::Ones({2, 3}, Dtype::kFloat32)},
            {Array::Full({2, 3}, 0.01f), Array::Full({3, 2}, 0.01f)});
}

TEST_P(ManipulationTest, Reshape) {
    using T = int32_t;
    Shape input_shape{2, 3, 4};
    Shape output_shape{3, 4, 2};

    Array a = testing::BuildArray(input_shape).WithLinearData<T>();
    Array b = Reshape(a, output_shape);
    ASSERT_EQ(output_shape, b.shape());
    EXPECT_EQ(a.data().get(), b.data().get()) << "Reshape must be done without copying data";
    Array e = testing::BuildArray(output_shape).WithLinearData<T>();
    testing::ExpectEqual(e, b);
}

TEST_P(ManipulationTest, ReshapeBugfix1) {
    using T = bool;
    Shape input_shape{6};
    Shape output_shape{2, 3};

    Array a = testing::BuildArray(input_shape).WithLinearData<T>();
    Array b = Reshape(a, output_shape);
    ASSERT_EQ(output_shape, b.shape());
    EXPECT_EQ(a.data().get(), b.data().get()) << "Reshape must be done without copying data";
    Array e = testing::BuildArray(output_shape).WithLinearData<T>();
    testing::ExpectEqual(e, b);
}

TEST_P(ManipulationTest, ReshapeBugfix2) {
    using T = double;
    Shape input_shape{2, 4};
    Shape output_shape{2, 1, 4, 1};

    Array a = testing::BuildArray(input_shape).WithLinearData<T>();
    Array b = Reshape(a, output_shape);
    ASSERT_EQ(output_shape, b.shape());
    EXPECT_EQ(a.data().get(), b.data().get()) << "Reshape must be done without copying data";
    Array e = testing::BuildArray(output_shape).WithLinearData<T>();
    testing::ExpectEqual(e, b);
}

// If an input array has a unit-length axis with 0-stride, that axis should not give rise to any copies.
TEST_P(ManipulationTest, ReshapeNoCopyZeroStrideAxis) {
    using T = int32_t;
    Shape input_shape_before_newaxis{2, 3, 4};
    Shape output_shape{3, 4, 2};

    // The shape of the input array is (2, 1, 3, 4) with strides (48, 0, 16, 4).
    Array a = (*testing::BuildArray(input_shape_before_newaxis).WithLinearData<T>()).At({Slice{}, NewAxis{}, Slice{}, Slice{}});
    assert(std::find(a.strides().begin(), a.strides().end(), 0) != a.strides().end());

    Array b = Reshape(a, output_shape);
    ASSERT_EQ(output_shape, b.shape());
    EXPECT_EQ(a.data().get(), b.data().get()) << "Reshape must be done without copying data";
    Array e = testing::BuildArray(output_shape).WithLinearData<T>();
    testing::ExpectEqual(e, b);
}
<commit_msg>rename tests to more descriptive names<commit_after>#include "xchainer/routines/manipulation.h"

#include <algorithm>
#include <cstdint>
#include <memory>
#include <string>
#include <vector>

#include <gtest/gtest.h>
#include <nonstd/optional.hpp>

#include "xchainer/array.h"
#include "xchainer/array_index.h"
#include "xchainer/check_backward.h"
#include "xchainer/device_id.h"
#include "xchainer/error.h"
#include "xchainer/shape.h"
#include "xchainer/strides.h"
#include "xchainer/testing/array.h"
#include "xchainer/testing/array_check.h"
#include "xchainer/testing/device_session.h"

namespace xchainer {
namespace {

class ManipulationTest : public ::testing::TestWithParam<std::string> {
protected:
    void SetUp() override {
        const std::string& backend_name = GetParam();
        device_session_.emplace(DeviceId{backend_name, 0});
    }

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

private:
    nonstd::optional<testing::DeviceSession> device_session_;
};

TEST_P(ManipulationTest, AsScalar) {
    using T = float;
    T value = 2.0f;
    Array a = testing::BuildArray<T>({1, 1, 1}, {value}).WithPadding(1);
    Scalar s = AsScalar(a);

    ASSERT_EQ(s.dtype(), TypeToDtype<T>);
    EXPECT_EQ(static_cast<T>(s), value);
}

TEST_P(ManipulationTest, AsScalarInvalidZeroElement) {
    Array a = testing::BuildArray<float>({0}, {});
    EXPECT_THROW(AsScalar(a), DimensionError);
}

TEST_P(ManipulationTest, AsScalarInvalidMoreThanOneElements) {
    Array a = testing::BuildArray<float>({2}, {1.0f, 2.0f});
    EXPECT_THROW(AsScalar(a), DimensionError);
}

TEST_P(ManipulationTest, Transpose) {
    Array a = testing::BuildArray({2, 3})         //
                      .WithLinearData<int32_t>()  //
                      .WithPadding(0);
    Array b = Transpose(a);

    EXPECT_EQ(Shape({3, 2}), b.shape());
    EXPECT_EQ(Strides({4, 12}), b.strides());

    Array e = testing::BuildArray({3, 2}).WithData<int32_t>({0, 3, 1, 4, 2, 5});
    testing::ExpectEqual(e, b);
}

TEST_P(ManipulationTest, TransposeNoncontiguous) {
    Array a = testing::BuildArray({2, 3})         //
                      .WithLinearData<int32_t>()  //
                      .WithPadding(1);
    Array b = Transpose(a);

    EXPECT_EQ(Shape({3, 2}), b.shape());

    Array e = testing::BuildArray({3, 2}).WithData<int32_t>({0, 3, 1, 4, 2, 5});
    testing::ExpectEqual(e, b);
}

TEST_P(ManipulationTest, TransposeBackward) {
    CheckBackwardComputation(
            [](const std::vector<Array>& xs) -> std::vector<Array> { return {Transpose(xs[0])}; },
            {Array::Zeros({2, 3}, Dtype::kFloat32).RequireGrad()},
            {testing::BuildArray({3, 2}, {1.f, 2.f, 3.f, 4.f, 5.f, 6.f})},
            {Array::Full({2, 3}, 1e-5f)});
}

TEST_P(ManipulationTest, TransposeDoubleBackward) {
    CheckDoubleBackwardComputation(
            [](const std::vector<Array>& xs) -> std::vector<Array> {
                auto t = Transpose(xs[0]);
                return {t * t};  // to make it nonlinear
            },
            {(*testing::BuildArray({2, 3}, {1.f, -1.f, 2.f, -2.f, 3.f, -3.f})).RequireGrad()},
            {Array::Ones({3, 2}, Dtype::kFloat32).RequireGrad()},
            {Array::Ones({2, 3}, Dtype::kFloat32)},
            {Array::Full({2, 3}, 0.01f), Array::Full({3, 2}, 0.01f)});
}

TEST_P(ManipulationTest, Reshape) {
    using T = int32_t;
    Shape input_shape{2, 3, 4};
    Shape output_shape{3, 4, 2};

    Array a = testing::BuildArray(input_shape).WithLinearData<T>();
    Array b = Reshape(a, output_shape);
    ASSERT_EQ(output_shape, b.shape());
    EXPECT_EQ(a.data().get(), b.data().get()) << "Reshape must be done without copying data";
    Array e = testing::BuildArray(output_shape).WithLinearData<T>();
    testing::ExpectEqual(e, b);
}

TEST_P(ManipulationTest, ReshapeBugfixWithStrideOne) {
    using T = bool;
    Shape input_shape{6};
    Shape output_shape{2, 3};

    Array a = testing::BuildArray(input_shape).WithLinearData<T>();
    Array b = Reshape(a, output_shape);
    ASSERT_EQ(output_shape, b.shape());
    EXPECT_EQ(a.data().get(), b.data().get()) << "Reshape must be done without copying data";
    Array e = testing::BuildArray(output_shape).WithLinearData<T>();
    testing::ExpectEqual(e, b);
}

TEST_P(ManipulationTest, ReshapeBugfixNewAxisAtEnd) {
    using T = double;
    Shape input_shape{2, 4};
    Shape output_shape{2, 1, 4, 1};

    Array a = testing::BuildArray(input_shape).WithLinearData<T>();
    Array b = Reshape(a, output_shape);
    ASSERT_EQ(output_shape, b.shape());
    EXPECT_EQ(a.data().get(), b.data().get()) << "Reshape must be done without copying data";
    Array e = testing::BuildArray(output_shape).WithLinearData<T>();
    testing::ExpectEqual(e, b);
}

// If an input array has a unit-length axis with 0-stride, that axis should not give rise to any copies.
TEST_P(ManipulationTest, ReshapeNoCopyZeroStrideAxis) {
    using T = int32_t;
    Shape input_shape_before_newaxis{2, 3, 4};
    Shape output_shape{3, 4, 2};

    // The shape of the input array is (2, 1, 3, 4) with strides (48, 0, 16, 4).
    Array a = (*testing::BuildArray(input_shape_before_newaxis).WithLinearData<T>()).At({Slice{}, NewAxis{}, Slice{}, Slice{}});
    assert(std::find(a.strides().begin(), a.strides().end(), 0) != a.strides().end());

    Array b = Reshape(a, output_shape);
    ASSERT_EQ(output_shape, b.shape());
    EXPECT_EQ(a.data().get(), b.data().get()) << "Reshape must be done without copying data";
    Array e = testing::BuildArray(output_shape).WithLinearData<T>();
    testing::ExpectEqual(e, b);
}

{'fraction_non_alphanumeric': 0.17320143884892086, 'fraction_numerical': 0.023741007194244605, 'mean_word_length': 3.781169389509888, 'pattern_counts': {'":': 0, '<': 79, '<?xml version=': 0, '>': 67, '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': '1772938', 'n_tokens_mistral': 4120, 'n_tokens_neox': 3918, 'n_words': 793}