Fix gradfn.ReduceMean.Backward not respecting operator shape

nlpodyssey · Nov 4, 2023 · 93ae88e · 93ae88e
1 parent e547f4b
commit 93ae88e
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 1 deletion.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -7,6 +7,11 @@ to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
 ## [Unreleased]
 
+### Fixed
+
+- `gradfn.ReduceMean.Backward` was not respecting the operand's shape, causing
+  incompatible-shape error when reduce-mean was applied to non-vector values.
+
 ## [1.1.0] - 2023-10-30
 
 ### Changed

diff --git a/mat/gradfn/reducemean.go b/mat/gradfn/reducemean.go
@@ -42,7 +42,7 @@ func (r *ReduceMean[O]) Backward(gy mat.Tensor) error {
 		x := r.x.Value().(mat.Matrix)
 		size := x.Size()
 		v := gy.Item().F64() / float64(size)
-		gx := x.NewMatrix(mat.WithShape(size), mat.WithBacking(mat.CreateInitializedSlice(size, v)))
+		gx := x.NewMatrix(mat.WithShape(x.Shape()...), mat.WithBacking(mat.CreateInitializedSlice(size, v)))
 		r.x.AccGrad(gx)
 	}
 	return nil