[BREAKING CHANGE] Adapt nonzero and where APIs to match NumPy

[ClaudiaComito]

This PR makes heat.nonzero and heat.where behave like the Numpy counterparts.

This is an API-breaking, non-backwards-compatible change.

See this example for the effect of this PR:

import heat as ht
import torch
import numpy as np

a = ht.arange(7)

print(torch.nonzero(a.larray > 3))
# tensor([[4],
#         [5],
#         [6]])

print(np.nonzero(a.numpy() > 3))
# (array([4, 5, 6]),)

print(ht.nonzero(a > 3))
# on main: DNDarray([4, 5, 6], dtype=ht.int64, device=cpu:0, split=None)
# on features/nonzero-updates: (DNDarray(MPI-rank: 0, Shape: (3,), Split: None, Local Shape: (3,), Device: cpu:0, Dtype: int64, Data:
#                                        [4, 5, 6]),)

# new on features/nonzero-updates:
print(ht.nonzero(a > 3, as_tuple=False))
# DNDarray([[4],
#           [5],
#           [6]], dtype=ht.int64, device=cpu:0, split=None)

[brownbaerchen]

Why is this needed? Seems like duplicate entries would be a bug at this point. Or are there some side effects of unique here?

[brownbaerchen]

I don't like that this is needed because it's not clean, but the tests don't pass without, so..

[JuanPedroGHM]

@brownbaerchen @ClaudiaComito I think this should be addressed before we proceed with the merge. Or has there been any updates. I would wait until this is resolved first before any merge.

[ClaudiaComito]

I use unique() here because I need the global vectorized sorting

[ClaudiaComito]

@JuanPedroGHM to me this is resolved. The indices need to be sorted by coordinate axis. i.e.

[[0, 0, 1],
 [0, 1, 3],
 [0, 1, 5],
 [1, 0, 0],
 [2, 0, 1]]

unique(axis) does exactly that. I don't know what it means that it's not clean.

[ClaudiaComito]

Expanding with more details. Suppose the global key is

key = ht.array([[3, 3, 3],
       [5, 4, 9],
       [5, 6, 0],
       [6, 3, 8],
       [1, 6, 6]], split=0)

We need to align the key with the actual local positions of the elements of the indexed array. The first step is the vectorized sorting (please if there is a more appropriate term for this let me know). I.e. we want to sort key by the rows:

key = ht.array([[1, 6, 6],
       [3, 3, 3],
       [5, 4, 9],
       [5, 6, 0],
       [6, 3, 8]], split=0)

unique(axis=0) sorts by the rows keeping elements of the same original row together.

sort(axis) will sort the elements within each axis (row or column in this case) so the connection between elements of the same row will be lost.

I agree unique is overkill, but I'm pretty sure Heat doesn't have a dedicated vectorized sorting functionality that works in distributed mode (we have an old issue that we could revive though #363 )

>>>  np.sort(a, axis=0)
array([[1, 3, 0],
       [3, 3, 3],
       [5, 4, 6],
       [5, 6, 8],
       [6, 6, 9]])

>>> np.sort(a, axis=1)
array([[3, 3, 3],
       [4, 5, 9],
       [0, 5, 6],
       [3, 6, 8],
       [1, 6, 6]])

[brownbaerchen]

What's not clean about this is:

• This is using a non-obvious side effect of unique
• This is wasting compute eliminating non-existent non-unique values

What we need is essentially:

>>> import numpy as np
>>> a = np.array([[0, 0, 1],
...  [0, 1, 3],
...  [0, 1, 5],
...  [1, 0, 0],
...  [2, 0, 1]])
>>> a[np.argsort(a[:, 0])]
array([[0, 0, 1],
       [0, 1, 3],
       [0, 1, 5],
       [1, 0, 0],
       [2, 0, 1]])
>>> np.unique(a, axis=0)
array([[0, 0, 1],
       [0, 1, 3],
       [0, 1, 5],
       [1, 0, 0],
       [2, 0, 1]])

Using argsort is clean because the name of the function implies what it's doing and all that it's doing. The reader need not wonder why this function is called because it's clear what it does.

Problem: Heat doesn't have argsort. It has sort which returns the indices, so we could do

>>> import heat as ht
>>> b = ht.array(a)
>>> b[ht.sort(b[:, 0])[1]]
DNDarray(MPI-rank: 0, Shape: (5, 3), Split: None, Local Shape: (5, 3), Device: cpu:0, Dtype: int64, Data:
         [[0, 0, 1],
          [0, 1, 3],
          [0, 1, 5],
          [1, 0, 0],
          [2, 0, 1]])

But doesn't work in parallel..

What to do? :D

[brownbaerchen]

The issues arise because the API of nonzero is totally changed in this PR. See the following example:

import heat as ht
import torch
import numpy as np

a = ht.arange(7)

print(torch.nonzero(a.larray > 3))
# tensor([[4],
#         [5],
#         [6]])

print(np.nonzero(a.numpy() > 3))
# (array([4, 5, 6]),)

print(ht.nonzero(a > 3))
# on main: DNDarray([4, 5, 6], dtype=ht.int64, device=cpu:0, split=None)
# on features/nonzero-updates: (DNDarray(MPI-rank: 0, Shape: (3,), Split: None, Local Shape: (3,), Device: cpu:0, Dtype: int64, Data:
#                                        [4, 5, 6]),)

# new on features/nonzero-updates:
print(ht.nonzero(a > 3, as_tuple=False))
# DNDarray([[4],
#           [5],
#           [6]], dtype=ht.int64, device=cpu:0, split=None)

That is to say the previous API was neither numpy, nor torch and this PR changes it to default to numpy but with the option to get the torch behavior.

This is a good thing in principle, but I don't like to silently change the API. So, yes it's annoying that the current main relies on the previous API in a bunch of places and doesn't work with only these changes, but also we should clearly indicate this as a breaking, API-changing change.

We have two options:

• Adapt all uses of nonzero in the current main to use the new and improved API, then merge this PR
• Don't merge this PR and merge the changes together with the huge advanced indexing PR Feature: NumPy-compliant distributed advanced indexing #938

Either way, we need to mention the API changes in the release notes. What do you think, @JuanPedroGHM, @mtar, @ClaudiaComito?

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 91.75%. Comparing base (5d42619) to head (181b994).
⚠️ Report is 5 commits behind head on main.
✅ All tests successful. No failed tests found.

Additional details and impacted files

@@           Coverage Diff           @@
##             main    #2332   +/-   ##
=======================================
  Coverage   91.74%   91.75%           
=======================================
  Files          87       87           
  Lines       14120    14135   +15     
=======================================
+ Hits        12955    12970   +15     
  Misses       1165     1165           

Flag	Coverage Δ
unit	91.75% <100.00%> (+<0.01%)	⬆️

Flags with carried forward coverage won't be shown. Click here to find out more.

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

[mtar]

Everything works as expected. I only have some small comments.

[mtar]

Can we parameterize the tests now that we are more using pytest?

[mtar]

I think the if conditions can be reduced, like putting the innermost condition to the top reduces the repeated isinstance calls for example.

Suggested change

	# ---- binary where(cond, x, y) branch ------------------------------------
	if cond.split is not None and (isinstance(x, DNDarray) or isinstance(y, DNDarray)):
	if (isinstance(x, DNDarray) and cond.split != x.split) or (
	isinstance(y, DNDarray) and cond.split != y.split
	):
	if len(y.shape) >= 1 and y.shape[0] > 1:
	# Only raise if the "other" array has a meaningful first dimension.
	if isinstance(y, DNDarray) and len(y.shape) >= 1 and y.shape[0] > 1:
	raise NotImplementedError("binary op not implemented for different split axes")
	if cond.split is not None and isinstance(y, DNDarray) and len(y.shape) >= 1 and y.shape[0] > 1:
	if (isinstance(x, DNDarray) and cond.split != x.split) or cond.split != y.split):
	raise NotImplementedError("binary op not implemented for different split axes")

[JuanPedroGHM]

In like the changes, but I would not merge this until we find the reason why the global nonzero tensor needs to be pruned of duplicates.

Is this somehow addressed in the other PR's that affect this two functions?