Fixed a bug in normalizationHEADmaster

3 files changed, 45 insertions, 4 deletions

diff --git a/modules/codec/c++/data_math.hpp b/modules/codec/c++/data_math.hpp
index f423fcb..86336ad 100644
--- a/modules/codec/c++/data_math.hpp
+++ b/modules/codec/c++/data_math.hpp
@@ -99,7 +99,7 @@ public:
 		return values_.at(index);
 	}
 
-	data<schema::Tensor<Inner, Dims...>, encode::Native> operator+(const data<schema::Tensor<Inner, Dims...>, encode::Native>& rhs) {
+	data<schema::Tensor<Inner, Dims...>, encode::Native> operator+(const data<schema::Tensor<Inner, Dims...>, encode::Native>& rhs) const {
 		data<schema::Tensor<Inner, Dims...>, encode::Native> c;
 
 		rank_iterator<Dims...>::in_fixed_bounds([&](const data<schema::FixedArray<schema::UInt64, sizeof...(Dims)>, encode::Native>& index) -> error_or<void>{
@@ -110,7 +110,7 @@ public:
 		return c;
 	}
 
-	data<schema::Tensor<Inner, Dims...>, encode::Native> operator-(const data<schema::Tensor<Inner, Dims...>, encode::Native>& rhs) {
+	data<schema::Tensor<Inner, Dims...>, encode::Native> operator-(const data<schema::Tensor<Inner, Dims...>, encode::Native>& rhs) const {
 		data<schema::Tensor<Inner, Dims...>, encode::Native> c;
 
 		rank_iterator<Dims...>::in_fixed_bounds([&](const data<schema::FixedArray<schema::UInt64, sizeof...(Dims)>, encode::Native>& index) -> error_or<void>{
@@ -131,6 +131,28 @@ public:
 
 		return native_change;
 	}
+
+	data<schema::Tensor<Inner,Dims...>, encode::Native> operator*(const data<schema::Scalar<Inner>, encode::Native>& scal) const {
+		data<schema::Tensor<Inner, Dims...>, encode::Native> c;
+
+		rank_iterator<Dims...>::in_fixed_bounds([&](const data<schema::FixedArray<schema::UInt64, sizeof...(Dims)>, encode::Native>& index) -> error_or<void>{
+			c.at(index) = at(index) * scal.at({});
+			return make_void();
+		});
+
+		return c;
+	}
+
+	data<schema::Tensor<Inner,Dims...>, encode::Native> operator/(const data<schema::Scalar<Inner>, encode::Native>& scal) const {
+		data<schema::Tensor<Inner, Dims...>, encode::Native> c;
+
+		rank_iterator<Dims...>::in_fixed_bounds([&](const data<schema::FixedArray<schema::UInt64, sizeof...(Dims)>, encode::Native>& index) -> error_or<void>{
+			c.at(index) = at(index) / scal.at({});
+			return make_void();
+		});
+
+		return c;
+	}
 };
 
 }
diff --git a/modules/codec/c++/math.hpp b/modules/codec/c++/math.hpp
index 67d8ed3..ca43dd0 100644
--- a/modules/codec/c++/math.hpp
+++ b/modules/codec/c++/math.hpp
@@ -35,8 +35,8 @@ data<schema::Vector<T,D>, Encoding> normalize(const data<schema::Vector<T,D>>& i
 	auto sqrt_inp_dot = sqrt<T>(inp_dot);
 
 	saw::data<schema::Vector<T,D>, Encoding> out;
-	out.at({{0u}}).set(out.at({{0u}}).get() / sqrt_inp_dot.at({}).get());
-	out.at({{1u}}).set(out.at({{1u}}).get() / sqrt_inp_dot.at({}).get());
+	out.at({{0u}}).set(input.at({{0u}}).get() / sqrt_inp_dot.at({}).get());
+	out.at({{1u}}).set(input.at({{1u}}).get() / sqrt_inp_dot.at({}).get());
 	return out;
 }
 }
diff --git a/modules/codec/tests/math.cpp b/modules/codec/tests/math.cpp
index 4fb012f..cc371dd 100644
--- a/modules/codec/tests/math.cpp
+++ b/modules/codec/tests/math.cpp
@@ -75,4 +75,23 @@ SAW_TEST("Math/Dot"){
 
 	SAW_EXPECT(c.at({}).get() == 3.0, std::string{"Unexpected value for dot product "} + std::to_string(c.at({}).get()) );
 }
+
+SAW_TEST("Math/Tensor Mult Scalar"){
+	using namespace saw;
+
+	data<sch::Vector<sch::Float64, 2u>> a;
+	{
+		a.at({{0u}}) = 2.0;
+		a.at({{1u}}) = 1.0;
+	}
+	data<sch::Vector<sch::Float64, 2u>> b;
+	{
+		b.at({{0u}}) = -1.0;
+		b.at({{1u}}) = 5.0;
+	}
+
+	auto c = math::dot(a,b);
+
+	SAW_EXPECT(c.at({}).get() == 3.0, std::string{"Unexpected value for dot product "} + std::to_string(c.at({}).get()) );
+}
 }