Repository: pydata/sparse

Issue:
IndexError on different shape arrays when computing ufunc
I have identified the following behavior during https://github.com/dask/dask/pull/4525:

```python
(Pdb) p divide(totals, ns, dtype=dtype)
<COO: shape=(2, 2, 2), dtype=float64, nnz=6, fill_value=0.0>
(Pdb) mu
<COO: shape=(1, 1, 1), dtype=float64, nnz=1, fill_value=0.0>
(Pdb) p divide(totals, ns, dtype=dtype) - mu
*** IndexError: index 1 is out of bounds for axis 0 with size 1
```

I'm not sure what's the expected behavior in such situation for sparse arrays, but a subtraction of arrays in shape (2,2,2) and (1,1,1) works with NumPy arrays, would you say this is a bug @hameerabbasi ?

A snippet to reproduce an traceback follow:

```python
import numpy as np
import sparse

a = sparse.COO.from_numpy(np.ones((2, 2, 2)))
b = sparse.COO.from_numpy(np.ones((1, 1, 1)))

a - b
```

```python
Traceback (most recent call last):
  File "sparse_indexerror.py", line 7, in <module>
    a - b
  File "/home/pentschev/.local/lib/python3.5/site-packages/numpy/lib/mixins.py", line 25, in func
    return ufunc(self, other)
  File "/home/pentschev/.local/lib/python3.5/site-packages/sparse/coo/core.py", line 1454, in __array_ufunc__
    result = elemwise(ufunc, *inputs, **kwargs)
  File "/home/pentschev/.local/lib/python3.5/site-packages/sparse/coo/umath.py", line 48, in elemwise
    return _Elemwise(func, *args, **kwargs).get_result()
  File "/home/pentschev/.local/lib/python3.5/site-packages/sparse/coo/umath.py", line 444, in get_result
    r = self._get_func_coords_data(mask)
  File "/home/pentschev/.local/lib/python3.5/site-packages/sparse/coo/umath.py", line 579, in _get_func_coords_data
    func_coords, func_data = _get_expanded_coords_data(func_coords, func_data, params, self.shape)
  File "/home/pentschev/.local/lib/python3.5/site-packages/sparse/coo/umath.py", line 263, in _get_expanded_coords_data
    expanded_data = data[all_idx[first_dim]]
IndexError: index 1 is out of bounds for axis 0 with size 1
```

Discussion:
codecov[bot]: # [Codecov](https://codecov.io/gh/pydata/sparse/pull/243?src=pr&el=h1) Report
> Merging [#243](https://codecov.io/gh/pydata/sparse/pull/243?src=pr&el=desc) into [master](https://codecov.io/gh/pydata/sparse/commit/1bdd13708c1319821f6f5e8476334528024ea667?src=pr&el=desc) will **decrease** coverage by `0.18%`.
> The diff coverage is `40%`.

[![Impacted file tree graph](https://codecov.io/gh/pydata/sparse/pull/243/graphs/tree.svg?width=650&token=H212u0Uxxw&height=150&src=pr)](https://codecov.io/gh/pydata/sparse/pull/243?src=pr&el=tree)

```diff
@@            Coverage Diff             @@
##           master     #243      +/-   ##
==========================================
- Coverage   97.68%   97.49%   -0.19%     
==========================================
  Files          10       10              
  Lines        1555     1559       +4     
==========================================
+ Hits         1519     1520       +1     
- Misses         36       39       +3
```


| [Impacted Files](https://codecov.io/gh/pydata/sparse/pull/243?src=pr&el=tree) | Coverage Δ | |
|---|---|---|
| [sparse/coo/umath.py](https://codecov.io/gh/pydata/sparse/pull/243/diff?src=pr&el=tree#diff-c3BhcnNlL2Nvby91bWF0aC5weQ==) | `95.96% <40%> (-1.17%)` | :arrow_down: |

------

[Continue to review full report at Codecov](https://codecov.io/gh/pydata/sparse/pull/243?src=pr&el=continue).
> **Legend** - [Click here to learn more](https://docs.codecov.io/docs/codecov-delta)
> `Δ = absolute <relative> (impact)`, `ø = not affected`, `? = missing data`
> Powered by [Codecov](https://codecov.io/gh/pydata/sparse/pull/243?src=pr&el=footer). Last update [1bdd137...c0b8f07](https://codecov.io/gh/pydata/sparse/pull/243?src=pr&el=lastupdated). Read the [comment docs](https://docs.codecov.io/docs/pull-request-comments).


Patch:
diff --git a/sparse/coo/umath.py b/sparse/coo/umath.py
index bfb613b..6db24f5 100644
--- a/sparse/coo/umath.py
+++ b/sparse/coo/umath.py
@@ -260,7 +260,13 @@ def _get_expanded_coords_data(coords, data, params, broadcast_shape):
     dim = 0
 
     expanded_coords = np.empty((len(broadcast_shape), all_idx.shape[1]), dtype=np.intp)
-    expanded_data = data[all_idx[first_dim]]
+
+    if first_dim != -1:
+        expanded_data = data[all_idx[first_dim]]
+    else:
+        expanded_coords = all_idx
+        expanded_data = np.repeat(data, np.prod(broadcast_shape, dtype=np.int64))
+        return np.asarray(expanded_coords), np.asarray(expanded_data)
 
     for d, p, l in zip(range(len(broadcast_shape)), params, broadcast_shape):
         if p:


Test patch:
diff --git a/sparse/tests/test_coo.py b/sparse/tests/test_coo.py
index e7854df..af639f6 100644
--- a/sparse/tests/test_coo.py
+++ b/sparse/tests/test_coo.py
@@ -613,12 +613,16 @@ def test_elemwise_trinary(func, shape):
                                            ((1, 5), (5, 1)),
                                            ((3, 1), (3, 4)),
                                            ((3, 1), (1, 4)),
-                                           ((1, 4), (3, 4))])
+                                           ((1, 4), (3, 4)),
+                                           ((2, 2, 2), (1, 1, 1))])
 def test_binary_broadcasting(func, shape1, shape2):
-    xs = sparse.random(shape1, density=0.5)
+    density1 = 1 if np.prod(shape1) == 1 else 0.5
+    density2 = 1 if np.prod(shape2) == 1 else 0.5
+
+    xs = sparse.random(shape1, density=density1)
     x = xs.todense()
 
-    ys = sparse.random(shape2, density=0.5)
+    ys = sparse.random(shape2, density=density2)
     y = ys.todense()
 
     expected = func(x, y)