Add more operator overloads in the new Tensor class (#3822)

This commit is contained in:
Kevin Sawatzky 2026-02-19 11:03:19 -06:00 committed by GitHub
parent 417343c920
commit efefdb17b1
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 47 additions and 1 deletions

View file

@ -762,6 +762,24 @@ public:
data_[i] += o.data_[i];
return *this;
}
Tensor& operator-=(const Tensor& o)
{
for (size_t i = 0; i < data_.size(); ++i)
data_[i] -= o.data_[i];
return *this;
}
Tensor& operator*=(const Tensor& o)
{
for (size_t i = 0; i < data_.size(); ++i)
data_[i] *= o.data_[i];
return *this;
}
Tensor& operator/=(const Tensor& o)
{
for (size_t i = 0; i < data_.size(); ++i)
data_[i] /= o.data_[i];
return *this;
}
Tensor operator+(const Tensor& o) const
{
@ -784,6 +802,13 @@ public:
r.data_[i] = data_[i] / o.data_[i];
return r;
}
Tensor operator*(const Tensor& o) const
{
Tensor r(shape_);
for (size_t i = 0; i < data_.size(); ++i)
r.data_[i] = data_[i] * o.data_[i];
return r;
}
Tensor operator+(T val) const
{

View file

@ -350,6 +350,9 @@ TEST_CASE("Tensor element-wise arithmetic")
c = a / b;
REQUIRE(c[0] == 0.25);
c = a * b;
REQUIRE(c[0] == 4.0);
}
TEST_CASE("Tensor scalar arithmetic")
@ -378,7 +381,7 @@ TEST_CASE("Tensor scalar arithmetic")
REQUIRE(b[0] == 11.0);
}
TEST_CASE("Tensor compound addition with tensor")
TEST_CASE("Tensor compound arithmetic with tensor")
{
Tensor<double> a({3}, 0.0);
Tensor<double> b({3}, 0.0);
@ -388,6 +391,24 @@ TEST_CASE("Tensor compound addition with tensor")
REQUIRE(a[0] == 11.0);
REQUIRE(a[1] == 22.0);
REQUIRE(a[2] == 33.0);
a = {1.0, 2.0, 3.0};
b -= a;
REQUIRE(b[0] == 9.0);
REQUIRE(b[1] == 18.0);
REQUIRE(b[2] == 27.0);
b = {10.0, 20.0, 30.0};
a *= b;
REQUIRE(a[0] == 10.0);
REQUIRE(a[1] == 40.0);
REQUIRE(a[2] == 90.0);
a = {1.0, 2.0, 3.0};
b /= a;
REQUIRE(b[0] == 10.0);
REQUIRE(b[1] == 10.0);
REQUIRE(b[2] == 10.0);
}
TEST_CASE("Tensor comparison operators")