cp-book

ecnerwala's competitive programming library

View the Project on GitHub ecnerwala/cp-book

🧪 tensor.test.cpp

View this file on GitHub · Last update: 2023-12-17 09:56:16-05:00

Depends on

Code

#include <catch2/catch_test_macros.hpp>

#include "tensor.hpp"

#include <string>

TEST_CASE("Tensor", "[tensor]") {
	using ten = tensor<std::string, 2>;
	ten a({2, 3});
	a[{0,0}] = "0";
	a[{0,1}] = "1";
	a[{0,2}] = "2";
	a[{1,0}] = "3";
	a[{1,1}] = "4";
	a[{1,2}] = "5";

	const ten const_a = a;
	ten b = a;
	REQUIRE(b[{0,0}] == "0");
	REQUIRE(b[{0,1}] == "1");
	REQUIRE(b[{0,2}] == "2");
	REQUIRE(b[{1,0}] == "3");
	REQUIRE(b[{1,1}] == "4");
	REQUIRE(b[{1,2}] == "5");

	// Bounds checked
	REQUIRE(b.at({0,0}) == "0");
	REQUIRE(b.at({0,1}) == "1");
	REQUIRE(b.at({0,2}) == "2");
	REQUIRE(b.at({1,0}) == "3");
	REQUIRE(b.at({1,1}) == "4");
	REQUIRE(b.at({1,2}) == "5");

	REQUIRE(*b[0][0] == "0");
	REQUIRE(*b[0][1] == "1");
	REQUIRE(*b[0][2] == "2");
	REQUIRE(*b[1][0] == "3");
	REQUIRE(*b[1][1] == "4");
	REQUIRE(*b[1][2] == "5");

	REQUIRE(*const_a[0][0] == "0");
	REQUIRE(*const_a[0][1] == "1");
	REQUIRE(*const_a[0][2] == "2");
	REQUIRE(*const_a[1][0] == "3");
	REQUIRE(*const_a[1][1] == "4");
	REQUIRE(*const_a[1][2] == "5");
}

#include <catch2/catch_test_macros.hpp>
#include <array>
#include <string>
#line 2 "src/tensor.test.cpp"

#line 2 "src/tensor.hpp"

#line 4 "src/tensor.hpp"

template <typename T, int NDIMS> struct tensor_view {
	static_assert(NDIMS >= 0, "NDIMS must be nonnegative");

protected:
	std::array<int, NDIMS> shape;
	std::array<int, NDIMS> strides;
	T* data;

	tensor_view(std::array<int, NDIMS> shape_, std::array<int, NDIMS> strides_, T* data_) : shape(shape_), strides(strides_), data(data_) {}

public:
	tensor_view() : shape{0}, strides{0}, data(nullptr) {}

protected:
	int flatten_index(std::array<int, NDIMS> idx) const {
		int res = 0;
		for (int i = 0; i < NDIMS; i++) { res += idx[i] * strides[i]; }
		return res;
	}
	int flatten_index_checked(std::array<int, NDIMS> idx) const {
		int res = 0;
		for (int i = 0; i < NDIMS; i++) {
			assert(0 <= idx[i] && idx[i] < shape[i]);
			res += idx[i] * strides[i];
		}
		return res;
	}

public:
	T& operator[] (std::array<int, NDIMS> idx) const {
#ifdef _GLIBCXX_DEBUG
		return data[flatten_index_checked(idx)];
#else
		return data[flatten_index(idx)];
#endif
	}
	T& at(std::array<int, NDIMS> idx) const {
		return data[flatten_index_checked(idx)];
	}

	template <int D = NDIMS>
	typename std::enable_if<(0 < D), tensor_view<T, NDIMS-1>>::type operator[] (int idx) const {
		std::array<int, NDIMS-1> nshape; std::copy(shape.begin()+1, shape.end(), nshape.begin());
		std::array<int, NDIMS-1> nstrides; std::copy(strides.begin()+1, strides.end(), nstrides.begin());
		T* ndata = data + (strides[0] * idx);
		return tensor_view<T, NDIMS-1>(nshape, nstrides, ndata);
	}
	template <int D = NDIMS>
	typename std::enable_if<(0 < D), tensor_view<T, NDIMS-1>>::type at(int idx) const {
		assert(0 <= idx && idx < shape[0]);
		return operator[](idx);
	}

	template <int D = NDIMS>
	typename std::enable_if<(0 == D), T&>::type operator * () const {
		return *data;
	}

	template <typename U, int D> friend struct tensor_view;
	template <typename U, int D> friend struct tensor;
};

template <typename T, int NDIMS> struct tensor {
	static_assert(NDIMS >= 0, "NDIMS must be nonnegative");

protected:
	std::array<int, NDIMS> shape;
	std::array<int, NDIMS> strides;
	int len;
	T* data;

public:
	tensor() : shape{0}, strides{0}, len(0), data(nullptr) {}

	explicit tensor(std::array<int, NDIMS> shape_, const T& t = T()) {
		shape = shape_;
		len = 1;
		for (int i = NDIMS-1; i >= 0; i--) {
			strides[i] = len;
			len *= shape[i];
		}
		data = new T[len];
		std::fill(data, data + len, t);
	}

	tensor(const tensor& o) : shape(o.shape), strides(o.strides), len(o.len), data(new T[len]) {
		for (int i = 0; i < len; i++) {
			data[i] = o.data[i];
		}
	}

	tensor& operator=(tensor&& o) noexcept {
		using std::swap;
		swap(shape, o.shape);
		swap(strides, o.strides);
		swap(len, o.len);
		swap(data, o.data);
		return *this;
	}
	tensor(tensor&& o) : tensor() {
		*this = std::move(o);
	}
	tensor& operator=(const tensor& o) {
		return *this = tensor(o);
	}
	~tensor() { delete[] data; }

	using view_t = tensor_view<T, NDIMS>;
	view_t view() {
		return tensor_view<T, NDIMS>(shape, strides, data);
	}
	operator view_t() {
		return view();
	}

	using const_view_t = tensor_view<const T, NDIMS>;
	const_view_t view() const {
		return tensor_view<const T, NDIMS>(shape, strides, data);
	}
	operator const_view_t() const {
		return view();
	}

	T& operator[] (std::array<int, NDIMS> idx) { return view()[idx]; }
	T& at(std::array<int, NDIMS> idx) { return view().at(idx); }
	const T& operator[] (std::array<int, NDIMS> idx) const { return view()[idx]; }
	const T& at(std::array<int, NDIMS> idx) const { return view().at(idx); }

	template <int D = NDIMS>
	typename std::enable_if<(0 < D), tensor_view<T, NDIMS-1>>::type operator[] (int idx) {
		return view()[idx];
	}
	template <int D = NDIMS>
	typename std::enable_if<(0 < D), tensor_view<T, NDIMS-1>>::type at(int idx) {
		return view().at(idx);
	}

	template <int D = NDIMS>
	typename std::enable_if<(0 < D), tensor_view<const T, NDIMS-1>>::type operator[] (int idx) const {
		return view()[idx];
	}
	template <int D = NDIMS>
	typename std::enable_if<(0 < D), tensor_view<const T, NDIMS-1>>::type at(int idx) const {
		return view().at(idx);
	}

	template <int D = NDIMS>
	typename std::enable_if<(0 == D), T&>::type operator * () {
		return *view();
	}
	template <int D = NDIMS>
	typename std::enable_if<(0 == D), const T&>::type operator * () const {
		return *view();
	}
};
#line 4 "src/tensor.test.cpp"

#line 6 "src/tensor.test.cpp"

TEST_CASE("Tensor", "[tensor]") {
	using ten = tensor<std::string, 2>;
	ten a({2, 3});
	a[{0,0}] = "0";
	a[{0,1}] = "1";
	a[{0,2}] = "2";
	a[{1,0}] = "3";
	a[{1,1}] = "4";
	a[{1,2}] = "5";

	const ten const_a = a;
	ten b = a;
	REQUIRE(b[{0,0}] == "0");
	REQUIRE(b[{0,1}] == "1");
	REQUIRE(b[{0,2}] == "2");
	REQUIRE(b[{1,0}] == "3");
	REQUIRE(b[{1,1}] == "4");
	REQUIRE(b[{1,2}] == "5");

	// Bounds checked
	REQUIRE(b.at({0,0}) == "0");
	REQUIRE(b.at({0,1}) == "1");
	REQUIRE(b.at({0,2}) == "2");
	REQUIRE(b.at({1,0}) == "3");
	REQUIRE(b.at({1,1}) == "4");
	REQUIRE(b.at({1,2}) == "5");

	REQUIRE(*b[0][0] == "0");
	REQUIRE(*b[0][1] == "1");
	REQUIRE(*b[0][2] == "2");
	REQUIRE(*b[1][0] == "3");
	REQUIRE(*b[1][1] == "4");
	REQUIRE(*b[1][2] == "5");

	REQUIRE(*const_a[0][0] == "0");
	REQUIRE(*const_a[0][1] == "1");
	REQUIRE(*const_a[0][2] == "2");
	REQUIRE(*const_a[1][0] == "3");
	REQUIRE(*const_a[1][1] == "4");
	REQUIRE(*const_a[1][2] == "5");
}

// clang-format off
// @formatter:off
#pragma GCC diagnostic push
#pragma GCC diagnostic ignored "-Wpragmas"
#pragma GCC diagnostic ignored "-Wunknown-warning-option"
#pragma GCC diagnostic ignored "-Wmisleading-indentation"
#pragma GCC diagnostic ignored "-Wmultistatement-macros"
#include <bits/stdc++.h>
#include <catch2/catch_test_macros.hpp>
// src/tensor.hpp
template<typename T,int NDIMS>struct tensor_view{
static_assert(NDIMS>=0,"NDIMS must be nonnegative");
protected:
std::array<int,NDIMS>shape;
std::array<int,NDIMS>strides;
T*data;
tensor_view(std::array<int,NDIMS>shape_,std::array<int,NDIMS>strides_,T*data_):shape(shape_),strides(strides_),data(data_){}
public:
tensor_view():shape{0},strides{0},data(nullptr){}
protected:
int flatten_index(std::array<int,NDIMS>idx)const{
int res=0;
for(int i=0;i<NDIMS;i++){res+=idx[i]*strides[i];}
return res;
}
int flatten_index_checked(std::array<int,NDIMS>idx)const{
int res=0;
for(int i=0;i<NDIMS;i++){
assert(0<=idx[i]&&idx[i]<shape[i]);
res+=idx[i]*strides[i];
}
return res;
}
public:
T&operator[](std::array<int,NDIMS>idx)const{
#ifdef _GLIBCXX_DEBUG
return data[flatten_index_checked(idx)];
#else
return data[flatten_index(idx)];
#endif
}
T&at(std::array<int,NDIMS>idx)const{
return data[flatten_index_checked(idx)];
}
template<int D=NDIMS>
typename std::enable_if<(0<D),tensor_view<T,NDIMS-1>>::type operator[](int idx)const{
std::array<int,NDIMS-1>nshape;std::copy(shape.begin()+1,shape.end(),nshape.begin());
std::array<int,NDIMS-1>nstrides;std::copy(strides.begin()+1,strides.end(),nstrides.begin());
T*ndata=data+(strides[0]*idx);
return tensor_view<T,NDIMS-1>(nshape,nstrides,ndata);
}
template<int D=NDIMS>
typename std::enable_if<(0<D),tensor_view<T,NDIMS-1>>::type at(int idx)const{
assert(0<=idx&&idx<shape[0]);
return operator[](idx);
}
template<int D=NDIMS>
typename std::enable_if<(0==D),T&>::type operator*()const{
return*data;
}
template<typename U,int D>friend struct tensor_view;
template<typename U,int D>friend struct tensor;
};
template<typename T,int NDIMS>struct tensor{
static_assert(NDIMS>=0,"NDIMS must be nonnegative");
protected:
std::array<int,NDIMS>shape;
std::array<int,NDIMS>strides;
int len;
T*data;
public:
tensor():shape{0},strides{0},len(0),data(nullptr){}
explicit tensor(std::array<int,NDIMS>shape_,const T&t=T()){
shape=shape_;
len=1;
for(int i=NDIMS-1;i>=0;i--){
strides[i]=len;
len*=shape[i];
}
data=new T[len];
std::fill(data,data+len,t);
}
tensor(const tensor&o):shape(o.shape),strides(o.strides),len(o.len),data(new T[len]){
for(int i=0;i<len;i++){
data[i]=o.data[i];
}
}
tensor&operator=(tensor&&o)noexcept{
using std::swap;
swap(shape,o.shape);
swap(strides,o.strides);
swap(len,o.len);
swap(data,o.data);
return*this;
}
tensor(tensor&&o):tensor(){
*this=std::move(o);
}
tensor&operator=(const tensor&o){
return*this=tensor(o);
}
~tensor(){delete[]data;}
using view_t=tensor_view<T,NDIMS>;
view_t view(){
return tensor_view<T,NDIMS>(shape,strides,data);
}
operator view_t(){
return view();
}
using const_view_t=tensor_view<const T,NDIMS>;
const_view_t view()const{
return tensor_view<const T,NDIMS>(shape,strides,data);
}
operator const_view_t()const{
return view();
}
T&operator[](std::array<int,NDIMS>idx){return view()[idx];}
T&at(std::array<int,NDIMS>idx){return view().at(idx);}
const T&operator[](std::array<int,NDIMS>idx)const{return view()[idx];}
const T&at(std::array<int,NDIMS>idx)const{return view().at(idx);}
template<int D=NDIMS>
typename std::enable_if<(0<D),tensor_view<T,NDIMS-1>>::type operator[](int idx){
return view()[idx];
}
template<int D=NDIMS>
typename std::enable_if<(0<D),tensor_view<T,NDIMS-1>>::type at(int idx){
return view().at(idx);
}
template<int D=NDIMS>
typename std::enable_if<(0<D),tensor_view<const T,NDIMS-1>>::type operator[](int idx)const{
return view()[idx];
}
template<int D=NDIMS>
typename std::enable_if<(0<D),tensor_view<const T,NDIMS-1>>::type at(int idx)const{
return view().at(idx);
}
template<int D=NDIMS>
typename std::enable_if<(0==D),T&>::type operator*(){
return*view();
}
template<int D=NDIMS>
typename std::enable_if<(0==D),const T&>::type operator*()const{
return*view();
}
};
// src/tensor.test.cpp
TEST_CASE("Tensor","[tensor]"){
using ten=tensor<std::string,2>;
ten a({2,3});
a[{0,0}]="0";
a[{0,1}]="1";
a[{0,2}]="2";
a[{1,0}]="3";
a[{1,1}]="4";
a[{1,2}]="5";
const ten const_a=a;
ten b=a;
REQUIRE(b[{0,0}]=="0");
REQUIRE(b[{0,1}]=="1");
REQUIRE(b[{0,2}]=="2");
REQUIRE(b[{1,0}]=="3");
REQUIRE(b[{1,1}]=="4");
REQUIRE(b[{1,2}]=="5");
REQUIRE(b.at({0,0})=="0");
REQUIRE(b.at({0,1})=="1");
REQUIRE(b.at({0,2})=="2");
REQUIRE(b.at({1,0})=="3");
REQUIRE(b.at({1,1})=="4");
REQUIRE(b.at({1,2})=="5");
REQUIRE(*b[0][0]=="0");
REQUIRE(*b[0][1]=="1");
REQUIRE(*b[0][2]=="2");
REQUIRE(*b[1][0]=="3");
REQUIRE(*b[1][1]=="4");
REQUIRE(*b[1][2]=="5");
REQUIRE(*const_a[0][0]=="0");
REQUIRE(*const_a[0][1]=="1");
REQUIRE(*const_a[0][2]=="2");
REQUIRE(*const_a[1][0]=="3");
REQUIRE(*const_a[1][1]=="4");
REQUIRE(*const_a[1][2]=="5");
}
#pragma GCC diagnostic pop
// clang-format on
// @formatter:on
Back to top page