indexing bug: wrong in-place updates for 4D arrays with paired integer indices

[abagusetty]

The following pattern with a 4D array and paired integer index arrays produces incorrect results. With NumPy and CuPy, this advanced indexing expression updates only the diagonal blocks hj[0, 0, :, :] and hj[1, 1, :, :], leaving the off-diagonal blocks unchanged, as defined by NumPy’s semantics. In dpnp, however, the same code corrupts the off-diagonal blocks as well.

reproducer:

import dpnp as np

natm = 2
hj = np.arange(2 * 2 * 3 * 3, dtype=float).reshape(2, 2, 3, 3)
tmp = -np.ones((natm, 3, 3), dtype=float)

print("hj before:\n", hj)
hj[range(natm), range(natm)] += 2.0 * tmp
print("hj after:\n", hj)

dpnp output:

hj before:
 [[[[ 0.  1.  2.]
   [ 3.  4.  5.]
   [ 6.  7.  8.]]

  [[ 9. 10. 11.]
   [12. 13. 14.]
   [15. 16. 17.]]]


 [[[18. 19. 20.]
   [21. 22. 23.]
   [24. 25. 26.]]

  [[27. 28. 29.]
   [30. 31. 32.]
   [33. 34. 35.]]]]
hj after:
 [[[[-2. -1.  0.]
   [ 1.  2.  3.]
   [ 4.  5.  6.]]

  [[ 7.  8.  9.]
   [10. 11. 12.]
   [13. 14. 15.]]]


 [[[16. 17. 18.]
   [19. 20. 21.]
   [22. 23. 24.]]

  [[25. 26. 27.]
   [28. 29. 30.]
   [31. 32. 33.]]]

vs (numpy/cupy) output:

hj before:
 [[[[ 0.  1.  2.]
   [ 3.  4.  5.]
   [ 6.  7.  8.]]

  [[ 9. 10. 11.]
   [12. 13. 14.]
   [15. 16. 17.]]]


 [[[18. 19. 20.]
   [21. 22. 23.]
   [24. 25. 26.]]

  [[27. 28. 29.]
   [30. 31. 32.]
   [33. 34. 35.]]]]
hj after:
 [[[[-2. -1.  0.]
   [ 1.  2.  3.]
   [ 4.  5.  6.]]

  [[ 9. 10. 11.]
   [12. 13. 14.]
   [15. 16. 17.]]]


 [[[18. 19. 20.]
   [21. 22. 23.]
   [24. 25. 26.]]

  [[25. 26. 27.]
   [28. 29. 30.]
   [31. 32. 33.]]]]

[colleeneb]

Any update by any chance?

[antonwolfy]

@colleeneb, I'm afraid no. Most probably there is some fix needed in dpctl, but that is locked due to ongoing migration pf the tensor to dpnp. We will resume work on that, once the migration is completed.

[ndgrigorian]

@abagusetty @colleeneb @vlad-perevezentsev
I coincidentally may have stumbled onto the answer to this while I was working on adding back C API functions to usm_ndarray.

In what was dpctl's _slicing.pxi, I noticed something suspicious.

cdef bint _is_boolean(object x) except *:
    """Gives True if x is an integral slice spec"""
    if isinstance(x, (ndarray, usm_ndarray)):
        if x.ndim > 0:
            return False
        if x.dtype.kind not in "b":
            return False
        return True
    if isinstance(x, bool):
        return True
    if isinstance(x, (int, float, complex)):
        return False
    if _is_buffer(x):
        mbuf = memoryview(x)
        if mbuf.ndim == 0:
            f = mbuf.format
            return f in "?"
        else:
            return False
    if callable(getattr(x, "__bool__", None)):
        try:
            x.__bool__()
        except (TypeError, ValueError):
            return False
        return True
    return False

note the getattr(x, "__bool__", None) branch. No other part of the _basic_slice_meta implementation could cause this, as most other iterables without __bool__ should pass through, as can be seen in the later error:

raise IndexError(
    "Only integers, slices (`:`), ellipsis (`...`), "
    "dpnp.tensor.newaxis (`None`) and integer and "
    "boolean arrays are valid indices."
)

added some print statements to prove it:

In [1]: import dpnp as np

In [2]: natm = 2
   ...: hj = np.arange(2 * 2 * 3 * 3, dtype=float).reshape(2, 2, 3, 3)
   ...: tmp = -np.ones((natm, 3, 3), dtype=float)
   ...:
   ...: print("hj before:\n", hj)
   ...: hj[range(natm), range(natm)] += 2.0 * tmp
   ...: print("hj after:\n", hj)
hj before:
 [[[[ 0.  1.  2.]
   [ 3.  4.  5.]
   [ 6.  7.  8.]]

  [[ 9. 10. 11.]
   [12. 13. 14.]
   [15. 16. 17.]]]


 [[[18. 19. 20.]
   [21. 22. 23.]
   [24. 25. 26.]]

  [[27. 28. 29.]
   [30. 31. 32.]
   [33. 34. 35.]]]]
We're here
Whoops!
We're here
Whoops!
We're here
Whoops!
We're here
Whoops!
We're here
Whoops!
We're here
Whoops!
We're here
Whoops!
We're here
Whoops!
hj after:
 [[[[-2. -1.  0.]
   [ 1.  2.  3.]
   [ 4.  5.  6.]]

  [[ 7.  8.  9.]
   [10. 11. 12.]
   [13. 14. 15.]]]


 [[[16. 17. 18.]
   [19. 20. 21.]
   [22. 23. 24.]]

  [[25. 26. 27.]
   [28. 29. 30.]
   [31. 32. 33.]]]]

the root cause is that range exposes __bool__ for some ungodly reason, unlike list, tuple, etc.

The solution is just to remove the logic for __bool__ which likely is meant to catch np.bool_ scalar inputs, and instead, do a hard check for np.bool_

[abagusetty]

@ndgrigorian Thanks for the update and perfect timing indeed

[vlad-perevezentsev]

Thanks @ndgrigorian for pointing out the problematic piece
Here are a few things after my investigation:

1. Removing the generic __bool__ fallback fixes the misclassification of range as a boolean scalar index. After that change range fails with an IndexError

IndexError: Only integers, slices (`:`), ellipsis (`...`), dpnp.tensor.newaxis (`None`) and integer and boolean arrays are valid indices.

which is consistent with the Array API indexing rules since range is not a supported index type.
so I think these changes make sense and can be kept.

2. dpnp should be compatible NumPy/CuPy so I will likely need to add a small compatibility layer that converts range/ list/ of integers into integer index arrays before passing them to dpnp.tensor to avoid an IndexError and add tests for that

I will open PR with these changes today

[abagusetty]

@vlad-perevezentsev Thanks for the PR. Just curious if this also applies to #2622

[vlad-perevezentsev]

@abagusetty originally I was working on #2783 but the fix #2872 naturally covers #2622 as well.
The _unwrap_index_element helper I added in dpnp_array.py converts list indices to numpy arrays via numpy.asarray which handles nested lists too

dpnp: 0.20.0dev6+26.ge64845ceac1d
dpctl: 0.23.0dev0+19.g86d7384e3b
frr: <class 'dpnp.dpnp_array.dpnp_array'> (3, 10240)

Case A: list-of-lists index: [[0, 1], [1, 2]] (type: <class 'list'> )
Case A: OK, out.shape = (2, 2, 10240)

Case B: NumPy array index dtype: int64 (type: <class 'numpy.ndarray'> )
Case B: OK, out.shape = (2, 2, 10240)

Case C: dpnp array index dtype: int64 (type: <class 'dpnp.dpnp_array.dpnp_array'> )
Case C: OK, out.shape = (2, 2, 10240)

[abagusetty]

@vlad-perevezentsev Thanks for the update. I feel luckily today!