fix: index access issue in argmin and argmax complex (#3176)

Co-authored-by: Ianna Osborne <[email protected]>
scikit-hep · Jul 19, 2024 · 4d5d49a · 4d5d49a
1 parent a9e48eb
commit 4d5d49a
Show file tree

Hide file tree

Showing 4 changed files with 154 additions and 12 deletions.
diff --git a/awkward-cpp/src/cpu-kernels/awkward_reduce_argmax_complex.cpp b/awkward-cpp/src/cpu-kernels/awkward_reduce_argmax_complex.cpp
@@ -17,9 +17,9 @@ ERROR awkward_reduce_argmax_complex(
  for (int64_t i = 0; i < lenparents; i++) {
  int64_t parent = parents[i];
  if (toptr[parent] == -1 ||
- (fromptr[i * 2] > fromptr[toptr[parent * 2]] ||
- (fromptr[i * 2] == fromptr[toptr[parent * 2]] &&
- fromptr[i * 2 + 1] > fromptr[toptr[parent * 2 + 1]]))) {
+ (fromptr[i * 2] > fromptr[toptr[parent] * 2] ||
+ (fromptr[i * 2] == fromptr[toptr[parent] * 2] &&
+ fromptr[i * 2 + 1] > fromptr[toptr[parent] * 2 + 1]))) {
  toptr[parent] = i;
  }
  }

diff --git a/awkward-cpp/src/cpu-kernels/awkward_reduce_argmin_complex.cpp b/awkward-cpp/src/cpu-kernels/awkward_reduce_argmin_complex.cpp
@@ -17,9 +17,9 @@ ERROR awkward_reduce_argmin_complex(
  for (int64_t i = 0; i < lenparents; i++) {
  int64_t parent = parents[i];
  if (toptr[parent] == -1 ||
- (fromptr[i * 2] < fromptr[toptr[parent * 2]] ||
- (fromptr[i * 2] == fromptr[toptr[parent * 2]] &&
- fromptr[i * 2 + 1] < fromptr[toptr[parent * 2 + 1]]))) {
+ (fromptr[i * 2] < fromptr[toptr[parent] * 2] ||
+ (fromptr[i * 2] == fromptr[toptr[parent] * 2] &&
+ fromptr[i * 2 + 1] < fromptr[toptr[parent] * 2 + 1]))) {
  toptr[parent] = i;
  }
  }

diff --git a/kernel-specification.yml b/kernel-specification.yml
@@ -4038,9 +4038,9 @@ kernels:
  toptr[k] = -1
  for i in range(lenparents):
  parent = parents[i]
- if (toptr[parent] == -1 or (fromptr[i * 2] > fromptr[toptr[parent * 2]] or
- (fromptr[i * 2] == fromptr[toptr[parent * 2]] and
- fromptr[i * 2 + 1] > fromptr[toptr[parent * 2 + 1]]))):
+ if (toptr[parent] == -1 or (fromptr[i * 2] > fromptr[toptr[parent] * 2] or
+ (fromptr[i * 2] == fromptr[toptr[parent] * 2] and
+ fromptr[i * 2 + 1] > fromptr[toptr[parent] * 2 + 1]))):
  toptr[parent] = i
  automatic-tests: false
 
@@ -4150,9 +4150,9 @@ kernels:
  toptr[k] = -1
  for i in range(lenparents):
  parent = parents[i]
- if (toptr[parent] == -1 or (fromptr[i * 2] < fromptr[toptr[parent * 2]] or
- (fromptr[i * 2] == fromptr[toptr[parent * 2]] and
- fromptr[i * 2 + 1] < fromptr[toptr[parent * 2 + 1]]))):
+ if (toptr[parent] == -1 or (fromptr[i * 2] < fromptr[toptr[parent] * 2] or
+ (fromptr[i * 2] == fromptr[toptr[parent] * 2] and
+ fromptr[i * 2 + 1] < fromptr[toptr[parent] * 2 + 1]))):
  toptr[parent] = i
  automatic-tests: false
 

diff --git a/kernel-test-data.json b/kernel-test-data.json
@@ -23992,6 +23992,148 @@
  }
  ]
  },
+ {
+ "name": "awkward_reduce_argmin_complex",
+ "status": true,
+ "tests": [
+ {
+ "error": false,
+ "message": "",
+ "inputs": {
+ "fromptr": [],
+ "lenparents": 0,
+ "outlength": 0,
+ "parents": []
+ },
+ "outputs": {
+ "toptr": []
+ }
+ },
+ {
+ "error": false,
+ "message": "",
+ "inputs": {
+ "fromptr": [0, 0],
+ "lenparents": 1,
+ "outlength": 1,
+ "parents": [0]
+ },
+ "outputs": {
+ "toptr": [0]
+ }
+ },
+ {
+ "error": false,
+ "message": "",
+ "inputs": {
+ "fromptr": [1, 0, 0, 1],
+ "lenparents": 2,
+ "outlength": 1,
+ "parents": [0, 0]
+ },
+ "outputs": {
+ "toptr": [1]
+ }
+ },
+ {
+ "error": false,
+ "message": "",
+ "inputs": {
+ "fromptr": [2, 2, 3, 3, 5, 5, 7, 7, 11, 11, 13, 13, 17, 17, 19, 19, 23, 23],
+ "lenparents": 9,
+ "outlength": 6,
+ "parents": [0, 0, 0, 2, 2, 3, 4, 4, 5]
+ },
+ "outputs": {
+ "toptr": [0, -1, 3, 5, 6, 8]
+ }
+ },
+ {
+ "error": false,
+ "message": "",
+ "inputs": {
+ "fromptr": [1, 1, 0, 1, 0, 0, 1, 0, 0, 1, 0, 1],
+ "lenparents": 6,
+ "outlength": 4,
+ "parents": [0, 0, 0, 2, 2, 3]
+ },
+ "outputs": {
+ "toptr": [2, -1, 4, 5]
+ }
+ }
+ ]
+ },
+ {
+ "name": "awkward_reduce_argmax_complex",
+ "status": true,
+ "tests": [
+ {
+ "error": false,
+ "message": "",
+ "inputs": {
+ "fromptr": [],
+ "lenparents": 0,
+ "outlength": 0,
+ "parents": []
+ },
+ "outputs": {
+ "toptr": []
+ }
+ },
+ {
+ "error": false,
+ "message": "",
+ "inputs": {
+ "fromptr": [0, 0],
+ "lenparents": 1,
+ "outlength": 1,
+ "parents": [0]
+ },
+ "outputs": {
+ "toptr": [0]
+ }
+ },
+ {
+ "error": false,
+ "message": "",
+ "inputs": {
+ "fromptr": [1, 0, 0, 1],
+ "lenparents": 2,
+ "outlength": 1,
+ "parents": [0, 0]
+ },
+ "outputs": {
+ "toptr": [0]
+ }
+ },
+ {
+ "error": false,
+ "message": "",
+ "inputs": {
+ "fromptr": [2, 2, 3, 3, 5, 5, 7, 7, 11, 11, 13, 13, 17, 17, 19, 19, 23, 23],
+ "lenparents": 9,
+ "outlength": 6,
+ "parents": [0, 0, 0, 2, 2, 3, 4, 4, 5]
+ },
+ "outputs": {
+ "toptr": [2, -1, 4, 5, 7, 8]
+ }
+ },
+ {
+ "error": false,
+ "message": "",
+ "inputs": {
+ "fromptr": [1, 1, 0, 1, 0, 0, 1, 1, 0, 0, 0, 1],
+ "lenparents": 6,
+ "outlength": 4,
+ "parents": [0, 0, 0, 2, 2, 3]
+ },
+ "outputs": {
+ "toptr": [0, -1, 3, 5]
+ }
+ }
+ ]
+ },
  {
  "name": "awkward_reduce_max",
  "status": true,