support negitave axis for argmax argmin

ledao · Sep 11, 2019 · f1efff5 · f1efff5
1 parent 9f0c3a8
commit f1efff5
Show file tree

Hide file tree

Showing 2 changed files with 47 additions and 0 deletions.
diff --git a/stats.go b/stats.go
@@ -278,6 +278,9 @@ func Max(a *Arrf, axis ...int) *Arrf {
 }
 
 func (a *Arrf) ArgMax(axis int) *Arrf {
+	if axis < 0 {
+		axis = axis + len(a.shape)
+	}
 	restAxis := make([]int, len(a.shape)-1)
 	ta := a.Copy()
 	for i, t := 0, 0; i < len(ta.shape); i++ {
@@ -329,6 +332,9 @@ func ArgMax(a *Arrf, axis int) *Arrf {
 
 //fixme has bug
 func (a *Arrf) ArgMin(axis int) *Arrf {
+	if axis < 0 {
+		axis = axis + len(a.shape)
+	}
 	restAxis := make([]int, len(a.shape)-1)
 	ta := a.Copy()
 	for i, t := 0, 0; i < len(ta.shape); i++ {

diff --git a/stats_test.go b/stats_test.go
@@ -82,17 +82,39 @@ func TestArgMax(t *testing.T) {
 		[0, 1, 1, 0]], got `, arr.ArgMax(0))
 	}
 
+	if arr.ArgMax(-3).NotEqual(Array([]float64{0, 0, 0, 1, 0, 1, 1, 1, 0, 1, 1, 0}, 3, 4)).AnyTrue() {
+		t.Error(`Expected
+		[[0, 0, 0, 1],
+		[0, 1, 1, 1],
+		[0, 1, 1, 0]], got `, arr.ArgMax(0))
+	}
+
+
 	if arr.ArgMax(1).NotEqual(Array([]float64{2, 0, 0, 2, 2, 1, 2, 2}, 2, 4)).AnyTrue() {
 		t.Error(`Expected
 		[[2, 0, 0, 2],
        		[2, 1, 2, 2]], got `, arr.ArgMax(1))
 	}
 
+	if arr.ArgMax(-2).NotEqual(Array([]float64{2, 0, 0, 2, 2, 1, 2, 2}, 2, 4)).AnyTrue() {
+		t.Error(`Expected
+		[[2, 0, 0, 2],
+       		[2, 1, 2, 2]], got `, arr.ArgMax(1))
+	}
+
+
 	if arr.ArgMax(2).NotEqual(Array([]float64{2, 2, 0, 3, 2, 2}, 2, 3)).AnyTrue() {
 		t.Error(`Expected
 		[[2, 2, 0],
        		[3, 2, 2]], got `, arr.ArgMax(2))
 	}
+
+	if arr.ArgMax(-1).NotEqual(Array([]float64{2, 2, 0, 3, 2, 2}, 2, 3)).AnyTrue() {
+		t.Error(`Expected
+		[[2, 2, 0],
+       		[3, 2, 2]], got `, arr.ArgMax(2))
+	}
+
 }
 
 func TestArgMin(t *testing.T) {
@@ -117,4 +139,23 @@ func TestArgMin(t *testing.T) {
 		[[3, 0, 1],
        		[2, 0, 0]], got `, arr.ArgMin(2))
 	}
+
+	if arr.ArgMin(-3).NotEqual(Array([]float64{1, 1, 1, 0, 1, 0, 0, 0, 1, 0, 0, 1}, 3, 4)).AnyTrue() {
+		t.Error(`Expected
+		[[1, 1, 1, 0],
+		[1, 0, 0, 0],
+		[1, 0, 0, 1]], got `, arr.ArgMin(0))
+	}
+
+	if arr.ArgMin(-2).NotEqual(Array([]float64{1, 2, 2, 0, 1, 0, 0, 1}, 2, 4)).AnyTrue() {
+		t.Error(`Expected
+		[[1, 2, 2, 0],
+       		[1, 0, 0, 1]], got `, arr.ArgMin(1))
+	}
+
+	if arr.ArgMin(-1).NotEqual(Array([]float64{3, 0, 1, 2, 0, 0}, 2, 3)).AnyTrue() {
+		t.Error(`Expected
+		[[3, 0, 1],
+       		[2, 0, 0]], got `, arr.ArgMin(2))
+	}
 }