Repository: pydata/sparse

Issue:
Reorganize coo.py?
`coo.py` is getting large and unweildly. I sometimes spend a lot of time scrolling through it or finding stuff in it. I was thinking of breaking it up into smaller chunks for better maintainability.

I propose the following:

- A `coo` directory.
- `coo/coo.py` (for the `COO` class)
- Move large implementations out of the `COO` class into methods, one for `elemwise`, one for `reduction`, one for `__getitem__`, and one for other methods.
- Similarly for `DOK`.

I love any proposed changes to this structure.

Discussion:
codecov-io: # [Codecov](https://codecov.io/gh/pydata/sparse/pull/116?src=pr&el=h1) Report
> Merging [#116](https://codecov.io/gh/pydata/sparse/pull/116?src=pr&el=desc) into [master](https://codecov.io/gh/pydata/sparse/commit/97d5690116f499f43e79c3debfe7d7881c8b10ba?src=pr&el=desc) will **decrease** coverage by `0.05%`.
> The diff coverage is `96.45%`.

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

```diff
@@            Coverage Diff             @@
##           master     #116      +/-   ##
==========================================
- Coverage   96.04%   95.98%   -0.06%     
==========================================
  Files          11       13       +2     
  Lines        1896     1920      +24     
==========================================
+ Hits         1821     1843      +22     
- Misses         75       77       +2
```

| Flag | Coverage Δ | |
|---|---|---|
| #python27 | `95.41% <95.83%> (+0.26%)` | :arrow_up: |
| #python36 | `95.78% <95.83%> (-0.27%)` | :arrow_down: |

| [Impacted Files](https://codecov.io/gh/pydata/sparse/pull/116?src=pr&el=tree) | Coverage Δ | |
|---|---|---|
| [sparse/tests/test\_coo.py](https://codecov.io/gh/pydata/sparse/pull/116/diff?src=pr&el=tree#diff-c3BhcnNlL3Rlc3RzL3Rlc3RfY29vLnB5) | `99.83% <100%> (ø)` | :arrow_up: |
| [sparse/coo/\_core.py](https://codecov.io/gh/pydata/sparse/pull/116/diff?src=pr&el=tree#diff-c3BhcnNlL2Nvby9fY29yZS5weQ==) | `93.69% <100%> (ø)` | |
| [sparse/coo/\_\_init\_\_.py](https://codecov.io/gh/pydata/sparse/pull/116/diff?src=pr&el=tree#diff-c3BhcnNlL2Nvby9fX2luaXRfXy5weQ==) | `100% <100%> (ø)` | |
| [sparse/compatibility.py](https://codecov.io/gh/pydata/sparse/pull/116/diff?src=pr&el=tree#diff-c3BhcnNlL2NvbXBhdGliaWxpdHkucHk=) | `84.61% <83.33%> (-15.39%)` | :arrow_down: |
| [sparse/coo/\_common.py](https://codecov.io/gh/pydata/sparse/pull/116/diff?src=pr&el=tree#diff-c3BhcnNlL2Nvby9fY29tbW9uLnB5) | `96.26% <96.26%> (ø)` | |
| [sparse/coo/\_elemwise.py](https://codecov.io/gh/pydata/sparse/pull/116/diff?src=pr&el=tree#diff-c3BhcnNlL2Nvby9fZWxlbXdpc2UucHk=) | `97.03% <97.03%> (ø)` | |

------

[Continue to review full report at Codecov](https://codecov.io/gh/pydata/sparse/pull/116?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/116?src=pr&el=footer). Last update [97d5690...74dae7e](https://codecov.io/gh/pydata/sparse/pull/116?src=pr&el=lastupdated). Read the [comment docs](https://docs.codecov.io/docs/pull-request-comments).

hameerabbasi: cc @mrocklin Feedback welcome. This is blocking some other work, so quick is good.
mrocklin: In general this seems fine to me.  I left a couple of small comments.
hameerabbasi: Hmm. Seems like there are some issues between Travis and GitHub today... On the other PR, Travis builds didn't report after finishing, and on this one, it isn't starting.

Patch:
diff --git a/.travis.yml b/.travis.yml
index ab81770..84462d0 100644
--- a/.travis.yml
+++ b/.travis.yml
@@ -25,11 +25,9 @@ install:
   # Install sparse
   - pip install -e .[tests-all]
 
-before_script:
-  - flake8
-
 script:
   - py.test --cov-config .coveragerc --cov=./
+  - flake8
 
 after_success:
   - codecov -F python${TRAVIS_PYTHON_VERSION//./}
diff --git a/docs/generated/sparse.COO.nanreduce.rst b/docs/generated/sparse.COO.nanreduce.rst
deleted file mode 100644
index 4f713b1..0000000
--- a/docs/generated/sparse.COO.nanreduce.rst
+++ /dev/null
@@ -1,6 +0,0 @@
-COO.nanreduce
-=============
-
-.. currentmodule:: sparse
-
-.. automethod:: COO.nanreduce
\ No newline at end of file
diff --git a/docs/generated/sparse.COO.rst b/docs/generated/sparse.COO.rst
index f66f9d3..1ca472e 100644
--- a/docs/generated/sparse.COO.rst
+++ b/docs/generated/sparse.COO.rst
@@ -44,8 +44,6 @@ COO
       COO.min
       COO.prod
 
-      COO.nanreduce
-
    .. rubric:: :ref:`Converting to other formats <converting>`
    .. autosummary::
       :toctree:
diff --git a/docs/generated/sparse.nanreduce.rst b/docs/generated/sparse.nanreduce.rst
new file mode 100644
index 0000000..454be0b
--- /dev/null
+++ b/docs/generated/sparse.nanreduce.rst
@@ -0,0 +1,6 @@
+nanreduce
+=========
+
+.. currentmodule:: sparse
+
+.. autofunction:: nanreduce
\ No newline at end of file
diff --git a/docs/generated/sparse.rst b/docs/generated/sparse.rst
index 0c4b825..451f65b 100644
--- a/docs/generated/sparse.rst
+++ b/docs/generated/sparse.rst
@@ -35,6 +35,8 @@ API
 
     nanprod
 
+    nanreduce
+
     nansum
 
     random
diff --git a/sparse/__init__.py b/sparse/__init__.py
index 488fdfa..fc11c47 100644
--- a/sparse/__init__.py
+++ b/sparse/__init__.py
@@ -1,10 +1,5 @@
-from .coo import COO, elemwise, tensordot, concatenate, stack, dot, triu, tril, where, \
-    nansum, nanprod, nanmin, nanmax
+from .coo import *
 from .dok import DOK
 from .sparse_array import SparseArray
 from .utils import random
 from ._version import __version__
-
-__all__ = ["SparseArray", "COO", "DOK",
-           "tensordot", "concatenate", "stack", "dot", "triu", "tril", "random", "where",
-           "nansum", "nanprod", "nanmin", "nanmax"]
diff --git a/sparse/compatibility.py b/sparse/compatibility.py
index f680bfe..77fdda6 100644
--- a/sparse/compatibility.py
+++ b/sparse/compatibility.py
@@ -1,21 +1,18 @@
 # flake8: noqa F401
+import sys
 
-try:
-    from itertools import izip_longest as zip_longest
-except ImportError:
+if sys.version_info < (2, 7):
+    raise ImportError('Need at least Python 2.7.')
+
+if sys.version_info[0] == 3 and sys.version_info[1] < 5:
+    raise ImportError('Need at least Python 3.5 if using Python 3.')
+
+if sys.version_info[0] >= 3:
     from itertools import zip_longest
+    from builtins import int, range, zip
+else:
+    from itertools import izip_longest as zip_longest
+    from itertools import izip as zip
 
-try:
     int = long
-except NameError:
-    from builtins import int
-
-try:
     range = xrange
-except NameError:
-    from builtins import range
-
-try:
-    from itertools import izip as zip
-except ImportError:
-    from builtins import zip
diff --git a/sparse/coo/__init__.py b/sparse/coo/__init__.py
new file mode 100644
index 0000000..c47dbe5
--- /dev/null
+++ b/sparse/coo/__init__.py
@@ -0,0 +1,4 @@
+from .core import COO
+from .umath import elemwise
+from .common import tensordot, dot, concatenate, stack, triu, tril, where, \
+    nansum, nanprod, nanmin, nanmax, nanreduce
diff --git a/sparse/coo/common.py b/sparse/coo/common.py
new file mode 100644
index 0000000..48baa90
--- /dev/null
+++ b/sparse/coo/common.py
@@ -0,0 +1,627 @@
+from functools import reduce
+import operator
+import warnings
+
+import numpy as np
+import scipy.sparse
+
+from ..sparse_array import SparseArray
+from ..compatibility import range
+from ..utils import isscalar
+
+
+def asCOO(x, name='asCOO', check=True):
+    """
+    Convert the input to :obj:`COO`. Passes through :obj:`COO` objects as-is.
+
+    Parameters
+    ----------
+    x : Union[SparseArray, scipy.sparse.spmatrix, numpy.ndarray]
+        The input array to convert.
+    name : str, optional
+        The name of the operation to use in the exception.
+    check : bool, optional
+        Whether to check for a dense input.
+
+    Returns
+    -------
+    COO
+        The converted :obj:`COO` array.
+
+    Raises
+    ------
+    ValueError
+        If ``check`` is true and a dense input is supplied.
+    """
+    from .core import COO
+
+    if check and not isinstance(x, (SparseArray, scipy.sparse.spmatrix)):
+        raise ValueError('Performing this operation would produce a dense result: %s' % name)
+
+    if not isinstance(x, COO):
+        x = COO(x)
+
+    return x
+
+
+def linear_loc(coords, shape, signed=False):
+    n = reduce(operator.mul, shape, 1)
+    if signed:
+        n = -n
+    dtype = np.min_scalar_type(n)
+    out = np.zeros(coords.shape[1], dtype=dtype)
+    tmp = np.zeros(coords.shape[1], dtype=dtype)
+    strides = 1
+    for i, d in enumerate(shape[::-1]):
+        np.multiply(coords[-(i + 1), :], strides, out=tmp, dtype=dtype)
+        np.add(tmp, out, out=out)
+        strides *= d
+    return out
+
+
+def tensordot(a, b, axes=2):
+    """
+    Perform the equivalent of :obj:`numpy.tensordot`.
+
+    Parameters
+    ----------
+    a, b : Union[COO, np.ndarray, scipy.sparse.spmatrix]
+        The arrays to perform the :code:`tensordot` operation on.
+    axes : tuple[Union[int, tuple[int], Union[int, tuple[int]], optional
+        The axes to match when performing the sum.
+
+    Returns
+    -------
+    Union[COO, numpy.ndarray]
+        The result of the operation.
+
+    See Also
+    --------
+    numpy.tensordot : NumPy equivalent function
+    """
+    # Much of this is stolen from numpy/core/numeric.py::tensordot
+    # Please see license at https://github.com/numpy/numpy/blob/master/LICENSE.txt
+    from .core import COO
+
+    try:
+        iter(axes)
+    except TypeError:
+        axes_a = list(range(-axes, 0))
+        axes_b = list(range(0, axes))
+    else:
+        axes_a, axes_b = axes
+    try:
+        na = len(axes_a)
+        axes_a = list(axes_a)
+    except TypeError:
+        axes_a = [axes_a]
+        na = 1
+    try:
+        nb = len(axes_b)
+        axes_b = list(axes_b)
+    except TypeError:
+        axes_b = [axes_b]
+        nb = 1
+
+    # a, b = asarray(a), asarray(b)  # <--- modified
+    as_ = a.shape
+    nda = a.ndim
+    bs = b.shape
+    ndb = b.ndim
+    equal = True
+    if na != nb:
+        equal = False
+    else:
+        for k in range(na):
+            if as_[axes_a[k]] != bs[axes_b[k]]:
+                equal = False
+                break
+            if axes_a[k] < 0:
+                axes_a[k] += nda
+            if axes_b[k] < 0:
+                axes_b[k] += ndb
+    if not equal:
+        raise ValueError("shape-mismatch for sum")
+
+    # Move the axes to sum over to the end of "a"
+    # and to the front of "b"
+    notin = [k for k in range(nda) if k not in axes_a]
+    newaxes_a = notin + axes_a
+    N2 = 1
+    for axis in axes_a:
+        N2 *= as_[axis]
+    newshape_a = (-1, N2)
+    olda = [as_[axis] for axis in notin]
+
+    notin = [k for k in range(ndb) if k not in axes_b]
+    newaxes_b = axes_b + notin
+    N2 = 1
+    for axis in axes_b:
+        N2 *= bs[axis]
+    newshape_b = (N2, -1)
+    oldb = [bs[axis] for axis in notin]
+
+    at = a.transpose(newaxes_a).reshape(newshape_a)
+    bt = b.transpose(newaxes_b).reshape(newshape_b)
+    res = _dot(at, bt)
+    if isinstance(res, scipy.sparse.spmatrix):
+        if res.nnz > reduce(operator.mul, res.shape) / 2:
+            res = res.todense()
+        else:
+            res = COO.from_scipy_sparse(res)  # <--- modified
+            res.has_duplicates = False
+    if isinstance(res, np.matrix):
+        res = np.asarray(res)
+    return res.reshape(olda + oldb)
+
+
+def dot(a, b):
+    """
+    Perform the equivalent of :obj:`numpy.dot` on two arrays.
+
+    Parameters
+    ----------
+    a, b : Union[COO, np.ndarray, scipy.sparse.spmatrix]
+        The arrays to perform the :code:`dot` operation on.
+
+    Returns
+    -------
+    Union[COO, numpy.ndarray]
+        The result of the operation.
+
+    See Also
+    --------
+    numpy.dot : NumPy equivalent function.
+    COO.dot : Equivalent function for COO objects.
+    """
+    if not hasattr(a, 'ndim') or not hasattr(b, 'ndim'):
+        raise NotImplementedError(
+            "Cannot perform dot product on types %s, %s" %
+            (type(a), type(b)))
+
+    if a.ndim == 1 and b.ndim == 1:
+        return (a * b).sum()
+
+    a_axis = -1
+    b_axis = -2
+
+    if b.ndim == 1:
+        b_axis = -1
+
+    return tensordot(a, b, axes=(a_axis, b_axis))
+
+
+def _dot(a, b):
+    from .core import COO
+
+    if isinstance(a, COO):
+        a.sum_duplicates()
+    if isinstance(b, COO):
+        b.sum_duplicates()
+    if isinstance(b, COO) and not isinstance(a, COO):
+        return _dot(b.T, a.T).T
+    aa = a.tocsr()
+
+    if isinstance(b, (COO, scipy.sparse.spmatrix)):
+        b = b.tocsc()
+    return aa.dot(b)
+
+
+def concatenate(arrays, axis=0):
+    """
+    Concatenate the input arrays along the given dimension.
+
+    Parameters
+    ----------
+    arrays : Iterable[Union[COO, numpy.ndarray, scipy.sparse.spmatrix]]
+        The input arrays to concatenate.
+    axis : int, optional
+        The axis along which to concatenate the input arrays. The default is zero.
+
+    Returns
+    -------
+    COO
+        The output concatenated array.
+
+    See Also
+    --------
+    numpy.concatenate : NumPy equivalent function
+    """
+    from .core import COO
+
+    arrays = [x if isinstance(x, COO) else COO(x) for x in arrays]
+    if axis < 0:
+        axis = axis + arrays[0].ndim
+    assert all(x.shape[ax] == arrays[0].shape[ax]
+               for x in arrays
+               for ax in set(range(arrays[0].ndim)) - {axis})
+    nnz = 0
+    dim = sum(x.shape[axis] for x in arrays)
+    shape = list(arrays[0].shape)
+    shape[axis] = dim
+
+    coords_dtype = np.min_scalar_type(max(shape) - 1) if len(shape) != 0 else np.uint8
+    data = np.concatenate([x.data for x in arrays])
+    coords = np.concatenate([x.coords for x in arrays], axis=1).astype(coords_dtype)
+
+    dim = 0
+    for x in arrays:
+        if dim:
+            coords[axis, nnz:x.nnz + nnz] += dim
+        dim += x.shape[axis]
+        nnz += x.nnz
+
+    has_duplicates = any(x.has_duplicates for x in arrays)
+
+    return COO(coords, data, shape=shape, has_duplicates=has_duplicates,
+               sorted=(axis == 0) and all(a.sorted for a in arrays))
+
+
+def stack(arrays, axis=0):
+    """
+    Stack the input arrays along the given dimension.
+
+    Parameters
+    ----------
+    arrays : Iterable[Union[COO, numpy.ndarray, scipy.sparse.spmatrix]]
+        The input arrays to stack.
+    axis : int, optional
+        The axis along which to stack the input arrays.
+
+    Returns
+    -------
+    COO
+        The output stacked array.
+
+    See Also
+    --------
+    numpy.stack : NumPy equivalent function
+    """
+    from .core import COO
+
+    assert len(set(x.shape for x in arrays)) == 1
+    arrays = [x if isinstance(x, COO) else COO(x) for x in arrays]
+    if axis < 0:
+        axis = axis + arrays[0].ndim + 1
+    data = np.concatenate([x.data for x in arrays])
+    coords = np.concatenate([x.coords for x in arrays], axis=1)
+    shape = list(arrays[0].shape)
+    shape.insert(axis, len(arrays))
+
+    coords_dtype = np.min_scalar_type(max(shape) - 1) if len(shape) != 0 else np.uint8
+
+    nnz = 0
+    dim = 0
+    new = np.empty(shape=(coords.shape[1],), dtype=coords_dtype)
+    for x in arrays:
+        new[nnz:x.nnz + nnz] = dim
+        dim += 1
+        nnz += x.nnz
+
+    has_duplicates = any(x.has_duplicates for x in arrays)
+    coords = [coords[i].astype(coords_dtype) for i in range(coords.shape[0])]
+    coords.insert(axis, new)
+    coords = np.stack(coords, axis=0)
+
+    return COO(coords, data, shape=shape, has_duplicates=has_duplicates,
+               sorted=(axis == 0) and all(a.sorted for a in arrays))
+
+
+def triu(x, k=0):
+    """
+    Returns an array with all elements below the k-th diagonal set to zero.
+
+    Parameters
+    ----------
+    x : COO
+        The input array.
+    k : int, optional
+        The diagonal below which elements are set to zero. The default is
+        zero, which corresponds to the main diagonal.
+
+    Returns
+    -------
+    COO
+        The output upper-triangular matrix.
+
+    See Also
+    --------
+    numpy.triu : NumPy equivalent function
+    """
+    from .core import COO
+
+    if not x.ndim >= 2:
+        raise NotImplementedError('sparse.triu is not implemented for scalars or 1-D arrays.')
+
+    mask = x.coords[-2] + k <= x.coords[-1]
+
+    coords = x.coords[:, mask]
+    data = x.data[mask]
+
+    return COO(coords, data, x.shape, x.has_duplicates, x.sorted)
+
+
+def tril(x, k=0):
+    """
+    Returns an array with all elements above the k-th diagonal set to zero.
+
+    Parameters
+    ----------
+    x : COO
+        The input array.
+    k : int, optional
+        The diagonal above which elements are set to zero. The default is
+        zero, which corresponds to the main diagonal.
+
+    Returns
+    -------
+    COO
+        The output lower-triangular matrix.
+
+    See Also
+    --------
+    numpy.tril : NumPy equivalent function
+    """
+    from .core import COO
+
+    if not x.ndim >= 2:
+        raise NotImplementedError('sparse.tril is not implemented for scalars or 1-D arrays.')
+
+    mask = x.coords[-2] + k >= x.coords[-1]
+
+    coords = x.coords[:, mask]
+    data = x.data[mask]
+
+    return COO(coords, data, x.shape, x.has_duplicates, x.sorted)
+
+
+def nansum(x, axis=None, keepdims=False, dtype=None, out=None):
+    """
+    Performs a ``NaN`` skipping sum operation along the given axes. Uses all axes by default.
+
+    Parameters
+    ----------
+    x : SparseArray
+        The array to perform the reduction on.
+    axis : Union[int, Iterable[int]], optional
+        The axes along which to sum. Uses all axes by default.
+    keepdims : bool, optional
+        Whether or not to keep the dimensions of the original array.
+    dtype: numpy.dtype
+        The data type of the output array.
+
+    Returns
+    -------
+    COO
+        The reduced output sparse array.
+
+    See Also
+    --------
+    :obj:`COO.sum` : Function without ``NaN`` skipping.
+    numpy.nansum : Equivalent Numpy function.
+    """
+    assert out is None
+    x = asCOO(x, name='nansum')
+    return nanreduce(x, np.add, axis=axis, keepdims=keepdims, dtype=dtype)
+
+
+def nanmax(x, axis=None, keepdims=False, dtype=None, out=None):
+    """
+    Maximize along the given axes, skipping ``NaN`` values. Uses all axes by default.
+
+    Parameters
+    ----------
+    x : SparseArray
+        The array to perform the reduction on.
+    axis : Union[int, Iterable[int]], optional
+        The axes along which to maximize. Uses all axes by default.
+    keepdims : bool, optional
+        Whether or not to keep the dimensions of the original array.
+    dtype: numpy.dtype
+        The data type of the output array.
+
+    Returns
+    -------
+    COO
+        The reduced output sparse array.
+
+    See Also
+    --------
+    :obj:`COO.max` : Function without ``NaN`` skipping.
+    numpy.nanmax : Equivalent Numpy function.
+    """
+    assert out is None
+    x = asCOO(x, name='nanmax')
+
+    ar = x.reduce(np.fmax, axis=axis, keepdims=keepdims,
+                  dtype=dtype)
+
+    if (isscalar(ar) and np.isnan(ar)) or np.isnan(ar.data).any():
+        warnings.warn("All-NaN slice encountered", RuntimeWarning, stacklevel=2)
+
+    return ar
+
+
+def nanmin(x, axis=None, keepdims=False, dtype=None, out=None):
+    """
+    Minimize along the given axes, skipping ``NaN`` values. Uses all axes by default.
+
+    Parameters
+    ----------
+    x : SparseArray
+        The array to perform the reduction on.
+    axis : Union[int, Iterable[int]], optional
+        The axes along which to minimize. Uses all axes by default.
+    keepdims : bool, optional
+        Whether or not to keep the dimensions of the original array.
+    dtype: numpy.dtype
+        The data type of the output array.
+
+    Returns
+    -------
+    COO
+        The reduced output sparse array.
+
+    See Also
+    --------
+    :obj:`COO.min` : Function without ``NaN`` skipping.
+    numpy.nanmin : Equivalent Numpy function.
+    """
+    assert out is None
+    x = asCOO(x, name='nanmin')
+
+    ar = x.reduce(np.fmin, axis=axis, keepdims=keepdims,
+                  dtype=dtype)
+
+    if (isscalar(ar) and np.isnan(ar)) or np.isnan(ar.data).any():
+        warnings.warn("All-NaN slice encountered", RuntimeWarning, stacklevel=2)
+
+    return ar
+
+
+def nanprod(x, axis=None, keepdims=False, dtype=None, out=None):
+    """
+    Performs a product operation along the given axes, skipping ``NaN`` values.
+    Uses all axes by default.
+
+    Parameters
+    ----------
+    x : SparseArray
+        The array to perform the reduction on.
+    axis : Union[int, Iterable[int]], optional
+        The axes along which to multiply. Uses all axes by default.
+    keepdims : bool, optional
+        Whether or not to keep the dimensions of the original array.
+    dtype: numpy.dtype
+        The data type of the output array.
+
+    Returns
+    -------
+    COO
+        The reduced output sparse array.
+
+    See Also
+    --------
+    :obj:`COO.prod` : Function without ``NaN`` skipping.
+    numpy.nanprod : Equivalent Numpy function.
+    """
+    assert out is None
+    x = asCOO(x)
+    return nanreduce(x, np.multiply, axis=axis, keepdims=keepdims, dtype=dtype)
+
+
+def where(condition, x=None, y=None):
+    """
+    Select values from either ``x`` or ``y`` depending on ``condition``.
+    If ``x`` and ``y`` are not given, returns indices where ``condition``
+    is nonzero.
+
+    Performs the equivalent of :obj:`numpy.where`.
+
+    Parameters
+    ----------
+    condition : SparseArray
+        The condition based on which to select values from
+        either ``x`` or ``y``.
+    x : SparseArray, optional
+        The array to select values from if ``condition`` is nonzero.
+    y : SparseArray, optional
+        The array to select values from if ``condition`` is zero.
+
+    Returns
+    -------
+    COO
+        The output array with selected values if ``x`` and ``y`` are given;
+        else where the array is nonzero.
+
+    Raises
+    ------
+    ValueError
+        If the operation would produce a dense result; or exactly one of
+        ``x`` and ``y`` are given.
+
+    See Also
+    --------
+    numpy.where : Equivalent Numpy function.
+    """
+    from .umath import elemwise
+
+    x_given = x is not None
+    y_given = y is not None
+
+    if not (x_given or y_given):
+        condition = asCOO(condition, name=str(np.where))
+        return tuple(condition.coords)
+
+    if x_given != y_given:
+        raise ValueError('either both or neither of x and y should be given')
+
+    return elemwise(np.where, condition, x, y)
+
+
+def _replace_nan(array, value):
+    """
+    Replaces ``NaN``s in ``array`` with ``value``.
+
+    Parameters
+    ----------
+    array : COO
+        The input array.
+    value : numpy.number
+        The values to replace ``NaN`` with.
+
+    Returns
+    -------
+    COO
+        A copy of ``array`` with the ``NaN``s replaced.
+    """
+    from .core import COO
+
+    if not np.issubdtype(array.dtype, np.floating):
+        return array
+
+    data = np.copy(array.data)
+    data[np.isnan(data)] = value
+
+    return COO(array.coords, data, shape=array.shape,
+               has_duplicates=array.has_duplicates,
+               sorted=array.sorted)
+
+
+def nanreduce(x, method, identity=None, axis=None, keepdims=False, **kwargs):
+    """
+    Performs an ``NaN`` skipping reduction on this array. See the documentation
+    on :obj:`COO.reduce` for examples.
+
+    Parameters
+    ----------
+    x : COO
+        The array to reduce.
+    method : numpy.ufunc
+        The method to use for performing the reduction.
+    identity : numpy.number
+        The identity value for this reduction. Inferred from ``method`` if not given.
+        Note that some ``ufunc`` objects don't have this, so it may be necessary to give it.
+    axis : Union[int, Iterable[int]], optional
+        The axes along which to perform the reduction. Uses all axes by default.
+    keepdims : bool, optional
+        Whether or not to keep the dimensions of the original array.
+    kwargs : dict
+        Any extra arguments to pass to the reduction operation.
+
+    Returns
+    -------
+    COO
+        The result of the reduction operation.
+
+    Raises
+    ------
+    ValueError
+        If reducing an all-zero axis would produce a nonzero result.
+
+    See Also
+    --------
+    COO.reduce : Similar method without ``NaN`` skipping functionality.
+    """
+    arr = _replace_nan(x, method.identity if identity is None else identity)
+    return arr.reduce(method, axis, keepdims, **kwargs)
diff --git a/sparse/coo.py b/sparse/coo/core.py
similarity index 59%
rename from sparse/coo.py
rename to sparse/coo/core.py
index 5100167..6fbfdd9 100644
--- a/sparse/coo.py
+++ b/sparse/coo/core.py
@@ -1,20 +1,16 @@
-from __future__ import absolute_import, division, print_function
-
-import warnings
-from collections import Iterable, defaultdict, deque
-from functools import reduce
-from itertools import product
 import numbers
-import operator
+from collections import Iterable, defaultdict, deque
 
 import numpy as np
 import scipy.sparse
 from numpy.lib.mixins import NDArrayOperatorsMixin
 
-from .slicing import normalize_index
-from .utils import _zero_of_dtype, isscalar, PositinalArgumentPartial
-from .sparse_array import SparseArray
-from .compatibility import int, zip_longest, range, zip
+from .common import dot
+from .umath import elemwise, broadcast_to
+from ..compatibility import int, range
+from ..slicing import normalize_index
+from ..sparse_array import SparseArray
+from ..utils import _zero_of_dtype
 
 
 class COO(SparseArray, NDArrayOperatorsMixin):
@@ -185,7 +181,7 @@ class COO(SparseArray, NDArrayOperatorsMixin):
         if cache:
             self.enable_caching()
         if data is None:
-            from .dok import DOK
+            from ..dok import DOK
 
             if isinstance(coords, COO):
                 self._make_shallow_copy_of(coords)
@@ -596,42 +592,6 @@ class COO(SparseArray, NDArrayOperatorsMixin):
 
         return self.reduce(method, **kwargs)
 
-    def nanreduce(self, method, identity=None, axis=None, keepdims=False, **kwargs):
-        """
-        Performs an ``NaN`` skipping reduction on this array. See the documentation
-        on :obj:`COO.reduce` for examples.
-
-        Parameters
-        ----------
-        method : numpy.ufunc
-            The method to use for performing the reduction.
-        identity : numpy.number
-            The identity value for this reduction. Inferred from ``method`` if not given.
-            Note that some ``ufunc`` objects don't have this, so it may be necessary to give it.
-        axis : Union[int, Iterable[int]], optional
-            The axes along which to perform the reduction. Uses all axes by default.
-        keepdims : bool, optional
-            Whether or not to keep the dimensions of the original array.
-        kwargs : dict
-            Any extra arguments to pass to the reduction operation.
-
-        Returns
-        -------
-        COO
-            The result of the reduction operation.
-
-        Raises
-        ------
-        ValueError
-            If reducing an all-zero axis would produce a nonzero result.
-
-        See Also
-        --------
-        COO.reduce : Similar method without ``NaN`` skipping functionality.
-        """
-        arr = _replace_nan(self, method.identity if identity is None else identity)
-        return arr.reduce(method, axis, keepdims, **kwargs)
-
     def reduce(self, method, axis=(0,), keepdims=False, **kwargs):
         """
         Performs a reduction operation on this array.
@@ -1231,7 +1191,9 @@ class COO(SparseArray, NDArrayOperatorsMixin):
         >>> np.array_equal(np.flatnonzero(x), s.linear_loc())
         True
         """
-        return _linear_loc(self.coords, self.shape, signed)
+        from .common import linear_loc
+
+        return linear_loc(self.coords, self.shape, signed)
 
     def reshape(self, shape):
         """
@@ -1514,15 +1476,7 @@ class COO(SparseArray, NDArrayOperatorsMixin):
         --------
         :obj:`numpy.broadcast_to` : NumPy equivalent function
         """
-        if shape == self.shape:
-            return self
-
-        result_shape = _get_broadcast_shape(self.shape, shape, is_result=True)
-        params = _get_broadcast_parameters(self.shape, result_shape)
-        coords, data = _get_expanded_coords_data(self.coords, self.data, params, result_shape)
-
-        return COO(coords, data, shape=result_shape, has_duplicates=self.has_duplicates,
-                   sorted=self.sorted)
+        return broadcast_to(self, shape)
 
     def round(self, decimals=0, out=None):
         """
@@ -1611,150 +1565,6 @@ class COO(SparseArray, NDArrayOperatorsMixin):
                              "large sparse array to dense")
 
 
-def tensordot(a, b, axes=2):
-    """
-    Perform the equivalent of :obj:`numpy.tensordot`.
-
-    Parameters
-    ----------
-    a, b : Union[COO, np.ndarray, scipy.sparse.spmatrix]
-        The arrays to perform the :code:`tensordot` operation on.
-    axes : tuple[Union[int, tuple[int], Union[int, tuple[int]], optional
-        The axes to match when performing the sum.
-
-    Returns
-    -------
-    Union[COO, numpy.ndarray]
-        The result of the operation.
-
-    See Also
-    --------
-    numpy.tensordot : NumPy equivalent function
-    """
-    # Much of this is stolen from numpy/core/numeric.py::tensordot
-    # Please see license at https://github.com/numpy/numpy/blob/master/LICENSE.txt
-    try:
-        iter(axes)
-    except TypeError:
-        axes_a = list(range(-axes, 0))
-        axes_b = list(range(0, axes))
-    else:
-        axes_a, axes_b = axes
-    try:
-        na = len(axes_a)
-        axes_a = list(axes_a)
-    except TypeError:
-        axes_a = [axes_a]
-        na = 1
-    try:
-        nb = len(axes_b)
-        axes_b = list(axes_b)
-    except TypeError:
-        axes_b = [axes_b]
-        nb = 1
-
-    # a, b = asarray(a), asarray(b)  # <--- modified
-    as_ = a.shape
-    nda = a.ndim
-    bs = b.shape
-    ndb = b.ndim
-    equal = True
-    if na != nb:
-        equal = False
-    else:
-        for k in range(na):
-            if as_[axes_a[k]] != bs[axes_b[k]]:
-                equal = False
-                break
-            if axes_a[k] < 0:
-                axes_a[k] += nda
-            if axes_b[k] < 0:
-                axes_b[k] += ndb
-    if not equal:
-        raise ValueError("shape-mismatch for sum")
-
-    # Move the axes to sum over to the end of "a"
-    # and to the front of "b"
-    notin = [k for k in range(nda) if k not in axes_a]
-    newaxes_a = notin + axes_a
-    N2 = 1
-    for axis in axes_a:
-        N2 *= as_[axis]
-    newshape_a = (-1, N2)
-    olda = [as_[axis] for axis in notin]
-
-    notin = [k for k in range(ndb) if k not in axes_b]
-    newaxes_b = axes_b + notin
-    N2 = 1
-    for axis in axes_b:
-        N2 *= bs[axis]
-    newshape_b = (N2, -1)
-    oldb = [bs[axis] for axis in notin]
-
-    at = a.transpose(newaxes_a).reshape(newshape_a)
-    bt = b.transpose(newaxes_b).reshape(newshape_b)
-    res = _dot(at, bt)
-    if isinstance(res, scipy.sparse.spmatrix):
-        if res.nnz > reduce(operator.mul, res.shape) / 2:
-            res = res.todense()
-        else:
-            res = COO.from_scipy_sparse(res)  # <--- modified
-            res.has_duplicates = False
-    if isinstance(res, np.matrix):
-        res = np.asarray(res)
-    return res.reshape(olda + oldb)
-
-
-def dot(a, b):
-    """
-    Perform the equivalent of :obj:`numpy.dot` on two arrays.
-
-    Parameters
-    ----------
-    a, b : Union[COO, np.ndarray, scipy.sparse.spmatrix]
-        The arrays to perform the :code:`dot` operation on.
-
-    Returns
-    -------
-    Union[COO, numpy.ndarray]
-        The result of the operation.
-
-    See Also
-    --------
-    numpy.dot : NumPy equivalent function.
-    COO.dot : Equivalent function for COO objects.
-    """
-    if not hasattr(a, 'ndim') or not hasattr(b, 'ndim'):
-        raise NotImplementedError(
-            "Cannot perform dot product on types %s, %s" %
-            (type(a), type(b)))
-
-    if a.ndim == 1 and b.ndim == 1:
-        return (a * b).sum()
-
-    a_axis = -1
-    b_axis = -2
-
-    if b.ndim == 1:
-        b_axis = -1
-
-    return tensordot(a, b, axes=(a_axis, b_axis))
-
-
-def _dot(a, b):
-    if isinstance(a, COO):
-        a.sum_duplicates()
-    if isinstance(b, COO):
-        b.sum_duplicates()
-    if isinstance(b, COO) and not isinstance(a, COO):
-        return _dot(b.T, a.T).T
-    aa = a.tocsr()
-
-    if isinstance(b, (COO, scipy.sparse.spmatrix)):
-        b = b.tocsc()
-    return aa.dot(b)
-
-
 def _keepdims(original, new, axis):
     shape = list(original.shape)
     for ax in axis:
@@ -1784,238 +1594,6 @@ def _mask(coords, idx, shape):
         return mask
 
 
-def _replace_nan(array, value):
-    """
-    Replaces ``NaN``s in ``array`` with ``value``.
-
-    Parameters
-    ----------
-    array : COO
-        The input array.
-    value : numpy.number
-        The values to replace ``NaN`` with.
-
-    Returns
-    -------
-    COO
-        A copy of ``array`` with the ``NaN``s replaced.
-    """
-    if not np.issubdtype(array.dtype, np.floating):
-        return array
-
-    data = np.copy(array.data)
-    data[np.isnan(data)] = value
-
-    return COO(array.coords, data, shape=array.shape,
-               has_duplicates=array.has_duplicates,
-               sorted=array.sorted)
-
-
-def concatenate(arrays, axis=0):
-    """
-    Concatenate the input arrays along the given dimension.
-
-    Parameters
-    ----------
-    arrays : Iterable[Union[COO, numpy.ndarray, scipy.sparse.spmatrix]]
-        The input arrays to concatenate.
-    axis : int, optional
-        The axis along which to concatenate the input arrays. The default is zero.
-
-    Returns
-    -------
-    COO
-        The output concatenated array.
-
-    See Also
-    --------
-    numpy.concatenate : NumPy equivalent function
-    """
-    arrays = [x if isinstance(x, COO) else COO(x) for x in arrays]
-    if axis < 0:
-        axis = axis + arrays[0].ndim
-    assert all(x.shape[ax] == arrays[0].shape[ax]
-               for x in arrays
-               for ax in set(range(arrays[0].ndim)) - {axis})
-    nnz = 0
-    dim = sum(x.shape[axis] for x in arrays)
-    shape = list(arrays[0].shape)
-    shape[axis] = dim
-
-    coords_dtype = np.min_scalar_type(max(shape) - 1) if len(shape) != 0 else np.uint8
-    data = np.concatenate([x.data for x in arrays])
-    coords = np.concatenate([x.coords for x in arrays], axis=1).astype(coords_dtype)
-
-    dim = 0
-    for x in arrays:
-        if dim:
-            coords[axis, nnz:x.nnz + nnz] += dim
-        dim += x.shape[axis]
-        nnz += x.nnz
-
-    has_duplicates = any(x.has_duplicates for x in arrays)
-
-    return COO(coords, data, shape=shape, has_duplicates=has_duplicates,
-               sorted=(axis == 0) and all(a.sorted for a in arrays))
-
-
-def stack(arrays, axis=0):
-    """
-    Stack the input arrays along the given dimension.
-
-    Parameters
-    ----------
-    arrays : Iterable[Union[COO, numpy.ndarray, scipy.sparse.spmatrix]]
-        The input arrays to stack.
-    axis : int, optional
-        The axis along which to stack the input arrays.
-
-    Returns
-    -------
-    COO
-        The output stacked array.
-
-    See Also
-    --------
-    numpy.stack : NumPy equivalent function
-    """
-    assert len(set(x.shape for x in arrays)) == 1
-    arrays = [x if isinstance(x, COO) else COO(x) for x in arrays]
-    if axis < 0:
-        axis = axis + arrays[0].ndim + 1
-    data = np.concatenate([x.data for x in arrays])
-    coords = np.concatenate([x.coords for x in arrays], axis=1)
-    shape = list(arrays[0].shape)
-    shape.insert(axis, len(arrays))
-
-    coords_dtype = np.min_scalar_type(max(shape) - 1) if len(shape) != 0 else np.uint8
-
-    nnz = 0
-    dim = 0
-    new = np.empty(shape=(coords.shape[1],), dtype=coords_dtype)
-    for x in arrays:
-        new[nnz:x.nnz + nnz] = dim
-        dim += 1
-        nnz += x.nnz
-
-    has_duplicates = any(x.has_duplicates for x in arrays)
-    coords = [coords[i].astype(coords_dtype) for i in range(coords.shape[0])]
-    coords.insert(axis, new)
-    coords = np.stack(coords, axis=0)
-
-    return COO(coords, data, shape=shape, has_duplicates=has_duplicates,
-               sorted=(axis == 0) and all(a.sorted for a in arrays))
-
-
-def triu(x, k=0):
-    """
-    Returns an array with all elements below the k-th diagonal set to zero.
-
-    Parameters
-    ----------
-    x : COO
-        The input array.
-    k : int, optional
-        The diagonal below which elements are set to zero. The default is
-        zero, which corresponds to the main diagonal.
-
-    Returns
-    -------
-    COO
-        The output upper-triangular matrix.
-
-    See Also
-    --------
-    numpy.triu : NumPy equivalent function
-    """
-    if not x.ndim >= 2:
-        raise NotImplementedError('sparse.triu is not implemented for scalars or 1-D arrays.')
-
-    mask = x.coords[-2] + k <= x.coords[-1]
-
-    coords = x.coords[:, mask]
-    data = x.data[mask]
-
-    return COO(coords, data, x.shape, x.has_duplicates, x.sorted)
-
-
-def tril(x, k=0):
-    """
-    Returns an array with all elements above the k-th diagonal set to zero.
-
-    Parameters
-    ----------
-    x : COO
-        The input array.
-    k : int, optional
-        The diagonal above which elements are set to zero. The default is
-        zero, which corresponds to the main diagonal.
-
-    Returns
-    -------
-    COO
-        The output lower-triangular matrix.
-
-    See Also
-    --------
-    numpy.tril : NumPy equivalent function
-    """
-    if not x.ndim >= 2:
-        raise NotImplementedError('sparse.tril is not implemented for scalars or 1-D arrays.')
-
-    mask = x.coords[-2] + k >= x.coords[-1]
-
-    coords = x.coords[:, mask]
-    data = x.data[mask]
-
-    return COO(coords, data, x.shape, x.has_duplicates, x.sorted)
-
-
-# (c) Paul Panzer
-# Taken from https://stackoverflow.com/a/47833496/774273
-# License: https://creativecommons.org/licenses/by-sa/3.0/
-def _match_arrays(a, b):
-    """
-    Finds all indexes into a and b such that a[i] = b[j]. The outputs are sorted
-    in lexographical order.
-
-    Parameters
-    ----------
-    a, b : np.ndarray
-        The input 1-D arrays to match. If matching of multiple fields is
-        needed, use np.recarrays. These two arrays must be sorted.
-
-    Returns
-    -------
-    a_idx, b_idx : np.ndarray
-        The output indices of every possible pair of matching elements.
-    """
-    if len(a) == 0 or len(b) == 0:
-        return np.array([], dtype=np.uint8), np.array([], dtype=np.uint8)
-    asw = np.r_[0, 1 + np.flatnonzero(a[:-1] != a[1:]), len(a)]
-    bsw = np.r_[0, 1 + np.flatnonzero(b[:-1] != b[1:]), len(b)]
-    al, bl = np.diff(asw), np.diff(bsw)
-    na = len(al)
-    asw, bsw = asw, bsw
-    abunq = np.r_[a[asw[:-1]], b[bsw[:-1]]]
-    m = np.argsort(abunq, kind='mergesort')
-    mv = abunq[m]
-    midx = np.flatnonzero(mv[:-1] == mv[1:])
-    ai, bi = m[midx], m[midx + 1] - na
-    aic = np.r_[0, np.cumsum(al[ai])]
-    a_idx = np.ones((aic[-1],), dtype=np.int_)
-    a_idx[aic[:-1]] = asw[ai]
-    a_idx[aic[1:-1]] -= asw[ai[:-1]] + al[ai[:-1]] - 1
-    a_idx = np.repeat(np.cumsum(a_idx), np.repeat(bl[bi], al[ai]))
-    bi = np.repeat(bi, al[ai])
-    bic = np.r_[0, np.cumsum(bl[bi])]
-    b_idx = np.ones((bic[-1],), dtype=np.int_)
-    b_idx[bic[:-1]] = bsw[bi]
-    b_idx[bic[1:-1]] -= bsw[bi[:-1]] + bl[bi[:-1]] - 1
-    b_idx = np.cumsum(b_idx)
-    return a_idx, b_idx
-
-
 def _grouped_reduce(x, groups, method, **kwargs):
     """
     Performs a :code:`ufunc` grouped reduce.
@@ -2049,769 +1627,3 @@ def _grouped_reduce(x, groups, method, **kwargs):
     result = method.reduceat(x, inv_idx, **kwargs)
     counts = np.diff(np.concatenate((inv_idx, [len(x)])))
     return result, inv_idx, counts
-
-
-def elemwise(func, *args, **kwargs):
-    """
-    Apply a function to any number of arguments.
-
-    Parameters
-    ----------
-    func : Callable
-        The function to apply. Must support broadcasting.
-    args : tuple, optional
-        The arguments to the function. Can be :obj:`SparseArray` objects
-        or :obj:`scipy.sparse.spmatrix` objects.
-    kwargs : dict, optional
-        Any additional arguments to pass to the function.
-
-    Returns
-    -------
-    COO
-        The result of applying the function.
-
-    Raises
-    ------
-    ValueError
-        If the operation would result in a dense matrix, or if the operands
-        don't have broadcastable shapes.
-
-    See Also
-    --------
-    :obj:`numpy.ufunc` : A similar Numpy construct. Note that any :code:`ufunc` can be used
-        as the :code:`func` input to this function.
-
-    Notes
-    -----
-    Previously, operations with Numpy arrays were sometimes supported. Now,
-    it is necessary to convert Numpy arrays to :obj:`COO` objects.
-    """
-    # Because we need to mutate args.
-    args = list(args)
-    posargs = []
-    pos = []
-    for i, arg in enumerate(args):
-        if isinstance(arg, scipy.sparse.spmatrix):
-            args[i] = COO.from_scipy_sparse(arg)
-        elif isscalar(arg) or (isinstance(arg, np.ndarray)
-                               and not arg.shape):
-            # Faster and more reliable to pass ()-shaped ndarrays as scalars.
-            if isinstance(arg, np.ndarray):
-                args[i] = arg[()]
-
-            pos.append(i)
-            posargs.append(args[i])
-        elif isinstance(arg, SparseArray) and not isinstance(arg, COO):
-            args[i] = COO(arg)
-        elif not isinstance(arg, COO):
-            raise ValueError("Performing this operation would produce "
-                             "a dense result: %s" % str(func))
-
-    # Filter out scalars as they are 'baked' into the function.
-    func = PositinalArgumentPartial(func, pos, posargs)
-    args = [arg for arg in args if not isscalar(arg)]
-
-    if len(args) == 0:
-        return func(**kwargs)
-
-    if len(args) == 1:
-        return _elemwise_unary(func, args[0], **kwargs)
-
-    return _elemwise_n_ary(func, *args, **kwargs)
-
-
-def _elemwise_n_ary(func, *args, **kwargs):
-    """
-    Apply a function to any number of arguments with broadcasting.
-
-    Parameters
-    ----------
-    func : Callable
-        The function to apply to arguments. Must support broadcasting.
-    args : list
-        Input :obj:`COO` or :obj:`numpy.ndarray`s.
-    kwargs : dict
-        Additional arguments to pass to the function.
-
-    Returns
-    -------
-    COO
-        The output array.
-
-    Raises
-    ------
-    ValueError
-        If the input shapes aren't compatible or the result will be dense.
-    """
-    args = list(args)
-
-    for arg in args:
-        if isinstance(arg, COO):
-            arg.sum_duplicates()
-
-    args_zeros = tuple(_zero_of_dtype(arg.dtype)[()] for arg in args)
-
-    func_value = func(*args_zeros, **kwargs)
-    func_zero = _zero_of_dtype(func_value.dtype)
-    if func_value != func_zero:
-        raise ValueError("Performing this operation would produce "
-                         "a dense result: %s" % str(func))
-    data_list = []
-    coords_list = []
-
-    cache = {}
-    for mask in product([True, False], repeat=len(args)):
-        if not any(mask):
-            continue
-
-        ci, di = _unmatch_coo(func, args, mask, cache, **kwargs)
-
-        coords_list.extend(ci)
-        data_list.extend(di)
-
-    result_shape = _get_nary_broadcast_shape(*[arg.shape for arg in args])
-
-    # Concatenate matches and mismatches
-    data = np.concatenate(data_list) if len(data_list) else np.empty((0,), dtype=args[0].dtype)
-    coords = np.concatenate(coords_list, axis=1) if len(coords_list) else \
-        np.empty((0, len(result_shape)), dtype=np.min_scalar_type(max(result_shape) - 1))
-
-    nonzero = data != func_zero
-    data = data[nonzero]
-    coords = coords[:, nonzero]
-
-    return COO(coords, data, shape=result_shape, has_duplicates=False)
-
-
-def _match_coo(*args, **kwargs):
-    """
-    Matches the coordinates for any number of input :obj:`COO` arrays.
-    Equivalent to "sparse" broadcasting for all arrays.
-
-    Parameters
-    ----------
-    args : Tuple[COO]
-        The input :obj:`COO` arrays.
-    return_midx : bool
-        Whether to return matched indices or matched arrays. Matching
-        only supported for two arrays. ``False`` by default.
-    cache : dict
-        Cache of things already matched. No cache by default.
-
-    Returns
-    -------
-    matched_idx : List[ndarray]
-        The indices of matched elements in the original arrays. Only returned if
-        ``return_midx`` is ``True``.
-    matched_arrays : List[COO]
-        The expanded, matched :obj:`COO` objects. Only returned if
-        ``return_midx`` is ``False``.
-    """
-    return_midx = kwargs.pop('return_midx', False)
-    cache = kwargs.pop('cache', None)
-
-    if kwargs:
-        raise ValueError('Unknown kwargs %s' % kwargs.keys())
-
-    if return_midx and (len(args) != 2 or cache is not None):
-        raise NotImplementedError('Matching indices only supported for two args, and no cache.')
-
-    matched_arrays = [args[0]]
-    cache_key = [id(args[0])]
-    for arg2 in args[1:]:
-        cache_key.append(id(arg2))
-        key = tuple(cache_key)
-        if cache is not None and key in cache:
-            matched_arrays = cache[key]
-            continue
-
-        cargs = [matched_arrays[0], arg2]
-        current_shape = _get_broadcast_shape(matched_arrays[0].shape, arg2.shape)
-        params = [_get_broadcast_parameters(arg.shape, current_shape) for arg in cargs]
-        reduced_params = [all(p) for p in zip(*params)]
-        reduced_shape = _get_reduced_shape(arg2.shape,
-                                           reduced_params[-arg2.ndim:])
-
-        reduced_coords = [_get_reduced_coords(arg.coords, reduced_params[-arg.ndim:])
-                          for arg in cargs]
-
-        linear = [_linear_loc(rc, reduced_shape) for rc in reduced_coords]
-        sorted_idx = [np.argsort(idx) for idx in linear]
-        linear = [idx[s] for idx, s in zip(linear, sorted_idx)]
-        matched_idx = _match_arrays(*linear)
-
-        if return_midx:
-            matched_idx = [sidx[midx] for sidx, midx in zip(sorted_idx, matched_idx)]
-            return matched_idx
-
-        coords = [arg.coords[:, s] for arg, s in zip(cargs, sorted_idx)]
-        mcoords = [c[:, idx] for c, idx in zip(coords, matched_idx)]
-        mcoords = _get_matching_coords(mcoords, params, current_shape)
-        mdata = [arg.data[sorted_idx[0]][matched_idx[0]] for arg in matched_arrays]
-        mdata.append(arg2.data[sorted_idx[1]][matched_idx[1]])
-        matched_arrays = [COO(mcoords, md, shape=current_shape) for md in mdata]
-
-        if cache is not None:
-            cache[key] = matched_arrays
-
-    return matched_arrays
-
-
-def _unmatch_coo(func, args, mask, cache, **kwargs):
-    """
-    Matches the coordinates for any number of input :obj:`COO` arrays.
-
-    First computes the matches, then filters out the non-matches.
-
-    Parameters
-    ----------
-    func : Callable
-        The function to compute matches
-    args : tuple[COO]
-        The input :obj:`COO` arrays.
-    mask : tuple[bool]
-        Specifies the inputs that are zero and the ones that are
-        nonzero.
-    kwargs: dict
-        Extra keyword arguments to pass to func.
-
-    Returns
-    -------
-    matched_coords : list[ndarray]
-        The matched coordinates.
-    matched_data : list[ndarray]
-        The matched data.
-    """
-    matched_args = [a for a, m in zip(args, mask) if m]
-    unmatched_args = [a for a, m in zip(args, mask) if not m]
-
-    matched_arrays = _match_coo(*matched_args, cache=cache)
-
-    pos = tuple(i for i, m in enumerate(mask) if not m)
-    posargs = [_zero_of_dtype(arg.dtype)[()] for arg, m in zip(args, mask) if not m]
-    result_shape = _get_nary_broadcast_shape(*[arg.shape for arg in args])
-
-    partial = PositinalArgumentPartial(func, pos, posargs)
-    matched_func = partial(*[a.data for a in matched_arrays], **kwargs)
-
-    unmatched_mask = matched_func != _zero_of_dtype(matched_func.dtype)
-
-    if not unmatched_mask.any():
-        return [], []
-
-    func_data = matched_func[unmatched_mask]
-    func_coords = matched_arrays[0].coords[:, unmatched_mask]
-
-    func_array = COO(func_coords, func_data, shape=matched_arrays[0].shape).broadcast_to(result_shape)
-
-    if all(mask):
-        return [func_array.coords], [func_array.data]
-
-    unmatched_mask = np.ones(func_array.nnz, dtype=np.bool)
-
-    for arg in unmatched_args:
-        matched_idx = _match_coo(func_array, arg, return_midx=True)[0]
-        unmatched_mask[matched_idx] = False
-
-    coords = np.asarray(func_array.coords[:, unmatched_mask], order='C')
-    data = np.asarray(func_array.data[unmatched_mask], order='C')
-
-    return [coords], [data]
-
-
-def _get_nary_broadcast_shape(*shapes):
-    """
-    Broadcast any number of shapes to a result shape.
-
-    Parameters
-    ----------
-    shapes : tuple[tuple[int]]
-        The shapes to broadcast.
-
-    Returns
-    -------
-    tuple[int]
-        The output shape.
-
-    Raises
-    ------
-    ValueError
-        If the input shapes cannot be broadcast to a single shape.
-    """
-    result_shape = ()
-
-    for shape in shapes:
-        try:
-            result_shape = _get_broadcast_shape(shape, result_shape)
-        except ValueError:
-            shapes_str = ', '.join(str(shape) for shape in shapes)
-            raise ValueError('operands could not be broadcast together with shapes %s'
-                             % shapes_str)
-
-    return result_shape
-
-
-def _get_broadcast_shape(shape1, shape2, is_result=False):
-    """
-    Get the overall broadcasted shape.
-
-    Parameters
-    ----------
-    shape1, shape2 : tuple[int]
-        The input shapes to broadcast together.
-    is_result : bool
-        Whether or not shape2 is also the result shape.
-
-    Returns
-    -------
-    result_shape : tuple[int]
-        The overall shape of the result.
-
-    Raises
-    ------
-    ValueError
-        If the two shapes cannot be broadcast together.
-    """
-    # https://stackoverflow.com/a/47244284/774273
-    if not all((l1 == l2) or (l1 == 1) or ((l2 == 1) and not is_result) for l1, l2 in
-               zip(shape1[::-1], shape2[::-1])):
-        raise ValueError('operands could not be broadcast together with shapes %s, %s' %
-                         (shape1, shape2))
-
-    result_shape = tuple(max(l1, l2) for l1, l2 in
-                         zip_longest(shape1[::-1], shape2[::-1], fillvalue=1))[::-1]
-
-    return result_shape
-
-
-def _get_broadcast_parameters(shape, broadcast_shape):
-    """
-    Get the broadcast parameters.
-
-    Parameters
-    ----------
-    shape : tuple[int]
-        The input shape.
-    broadcast_shape
-        The shape to broadcast to.
-
-    Returns
-    -------
-    params : list
-        A list containing None if the dimension isn't in the original array, False if
-        it needs to be broadcast, and True if it doesn't.
-    """
-    params = [None if l1 is None else l1 == l2 for l1, l2
-              in zip_longest(shape[::-1], broadcast_shape[::-1], fillvalue=None)][::-1]
-
-    return params
-
-
-def _get_reduced_coords(coords, params):
-    """
-    Gets only those dimensions of the coordinates that don't need to be broadcast.
-
-    Parameters
-    ----------
-    coords : np.ndarray
-        The coordinates to reduce.
-    params : list
-        The params from which to check which dimensions to get.
-
-    Returns
-    -------
-    reduced_coords : np.ndarray
-        The reduced coordinates.
-    """
-
-    reduced_params = [bool(param) for param in params]
-
-    return coords[reduced_params]
-
-
-def _get_reduced_shape(shape, params):
-    """
-    Gets only those dimensions of the coordinates that don't need to be broadcast.
-
-    Parameters
-    ----------
-    coords : np.ndarray
-        The coordinates to reduce.
-    params : list
-        The params from which to check which dimensions to get.
-
-    Returns
-    -------
-    reduced_coords : np.ndarray
-        The reduced coordinates.
-    """
-    reduced_shape = tuple(l for l, p in zip(shape, params) if p)
-
-    return reduced_shape
-
-
-def _get_expanded_coords_data(coords, data, params, broadcast_shape):
-    """
-    Expand coordinates/data to broadcast_shape. Does most of the heavy lifting for broadcast_to.
-    Produces sorted output for sorted inputs.
-
-    Parameters
-    ----------
-    coords : np.ndarray
-        The coordinates to expand.
-    data : np.ndarray
-        The data corresponding to the coordinates.
-    params : list
-        The broadcast parameters.
-    broadcast_shape : tuple[int]
-        The shape to broadcast to.
-
-    Returns
-    -------
-    expanded_coords : np.ndarray
-        List of 1-D arrays. Each item in the list has one dimension of coordinates.
-    expanded_data : np.ndarray
-        The data corresponding to expanded_coords.
-    """
-    first_dim = -1
-    expand_shapes = []
-    for d, p, l in zip(range(len(broadcast_shape)), params, broadcast_shape):
-        if p and first_dim == -1:
-            expand_shapes.append(coords.shape[1])
-            first_dim = d
-
-        if not p:
-            expand_shapes.append(l)
-
-    all_idx = _cartesian_product(*(np.arange(d, dtype=np.min_scalar_type(d - 1)) for d in expand_shapes))
-    dt = np.result_type(*(np.min_scalar_type(l - 1) for l in broadcast_shape))
-
-    false_dim = 0
-    dim = 0
-
-    expanded_coords = np.empty((len(broadcast_shape), all_idx.shape[1]), dtype=dt)
-    expanded_data = data[all_idx[first_dim]]
-
-    for d, p, l in zip(range(len(broadcast_shape)), params, broadcast_shape):
-        if p:
-            expanded_coords[d] = coords[dim, all_idx[first_dim]]
-        else:
-            expanded_coords[d] = all_idx[false_dim + (d > first_dim)]
-            false_dim += 1
-
-        if p is not None:
-            dim += 1
-
-    return np.asarray(expanded_coords), np.asarray(expanded_data)
-
-
-# (c) senderle
-# Taken from https://stackoverflow.com/a/11146645/774273
-# License: https://creativecommons.org/licenses/by-sa/3.0/
-def _cartesian_product(*arrays):
-    """
-    Get the cartesian product of a number of arrays.
-
-    Parameters
-    ----------
-    arrays : Tuple[np.ndarray]
-        The arrays to get a cartesian product of. Always sorted with respect
-        to the original array.
-
-    Returns
-    -------
-    out : np.ndarray
-        The overall cartesian product of all the input arrays.
-    """
-    broadcastable = np.ix_(*arrays)
-    broadcasted = np.broadcast_arrays(*broadcastable)
-    rows, cols = np.prod(broadcasted[0].shape), len(broadcasted)
-    dtype = np.result_type(*arrays)
-    out = np.empty(rows * cols, dtype=dtype)
-    start, end = 0, rows
-    for a in broadcasted:
-        out[start:end] = a.reshape(-1)
-        start, end = end, end + rows
-    return out.reshape(cols, rows)
-
-
-def _elemwise_unary(func, self, *args, **kwargs):
-    check = kwargs.pop('check', True)
-    data_zero = _zero_of_dtype(self.dtype)
-    zero_func = func(data_zero, *args, **kwargs)
-    func_zero = _zero_of_dtype(zero_func.dtype)
-    if check and zero_func != func_zero:
-        raise ValueError("Performing this operation would produce "
-                         "a dense result: %s" % str(func))
-
-    data_func = func(self.data, *args, **kwargs)
-    nonzero = data_func != func_zero
-
-    return COO(self.coords[:, nonzero], data_func[nonzero],
-               shape=self.shape,
-               has_duplicates=self.has_duplicates,
-               sorted=self.sorted)
-
-
-def _get_matching_coords(coords, params, shape):
-    """
-    Get the matching coords across a number of broadcast operands.
-
-    Parameters
-    ----------
-    coords : list[numpy.ndarray]
-        The input coordinates.
-    params : list[Union[bool, none]]
-        The broadcast parameters.
-    Returns
-    -------
-    numpy.ndarray
-        The broacasted coordinates
-    """
-    matching_coords = []
-    dims = np.zeros(len(coords), dtype=np.uint8)
-
-    for p_all in zip(*params):
-        for i, p in enumerate(p_all):
-            if p:
-                matching_coords.append(coords[i][dims[i]])
-                break
-        else:
-            matching_coords.append(coords[dims[0]])
-
-        for i, p in enumerate(p_all):
-            if p is not None:
-                dims[i] += 1
-
-    dtype = np.min_scalar_type(max(shape) - 1)
-
-    return np.asarray(matching_coords, dtype=dtype)
-
-
-def _linear_loc(coords, shape, signed=False):
-    n = reduce(operator.mul, shape, 1)
-    if signed:
-        n = -n
-    dtype = np.min_scalar_type(n)
-    out = np.zeros(coords.shape[1], dtype=dtype)
-    tmp = np.zeros(coords.shape[1], dtype=dtype)
-    strides = 1
-    for i, d in enumerate(shape[::-1]):
-        np.multiply(coords[-(i + 1), :], strides, out=tmp, dtype=dtype)
-        np.add(tmp, out, out=out)
-        strides *= d
-    return out
-
-
-def asCOO(x, name='asCOO', check=True):
-    """
-    Convert the input to :obj:`COO`. Passes through :obj:`COO` objects as-is.
-
-    Parameters
-    ----------
-    x : Union[SparseArray, scipy.sparse.spmatrix, numpy.ndarray]
-        The input array to convert.
-    name : str, optional
-        The name of the operation to use in the exception.
-    check : bool, optional
-        Whether to check for a dense input.
-
-    Returns
-    -------
-    COO
-        The converted :obj:`COO` array.
-
-    Raises
-    ------
-    ValueError
-        If ``check`` is true and a dense input is supplied.
-    """
-    if check and not isinstance(x, (SparseArray, scipy.sparse.spmatrix)):
-        raise ValueError('Performing this operation would produce a dense result: %s' % name)
-
-    if not isinstance(x, COO):
-        x = COO(x)
-
-    return x
-
-
-def where(condition, x=None, y=None):
-    """
-    Select values from either ``x`` or ``y`` depending on ``condition``.
-    If ``x`` and ``y`` are not given, returns indices where ``condition``
-    is nonzero.
-
-    Performs the equivalent of :obj:`numpy.where`.
-
-    Parameters
-    ----------
-    condition : SparseArray
-        The condition based on which to select values from
-        either ``x`` or ``y``.
-    x : SparseArray, optional
-        The array to select values from if ``condition`` is nonzero.
-    y : SparseArray, optional
-        The array to select values from if ``condition`` is zero.
-
-    Returns
-    -------
-    COO
-        The output array with selected values if ``x`` and ``y`` are given;
-        else where the array is nonzero.
-
-    Raises
-    ------
-    ValueError
-        If the operation would produce a dense result; or exactly one of
-        ``x`` and ``y`` are given.
-
-    See Also
-    --------
-    numpy.where : Equivalent Numpy function.
-    """
-    x_given = x is not None
-    y_given = y is not None
-
-    if not (x_given or y_given):
-        condition = asCOO(condition, name=str(np.where))
-        return tuple(condition.coords)
-
-    if x_given != y_given:
-        raise ValueError('either both or neither of x and y should be given')
-
-    return elemwise(np.where, condition, x, y)
-
-
-def nansum(x, axis=None, keepdims=False, dtype=None, out=None):
-    """
-    Performs a ``NaN`` skipping sum operation along the given axes. Uses all axes by default.
-
-    Parameters
-    ----------
-    x : SparseArray
-        The array to perform the reduction on.
-    axis : Union[int, Iterable[int]], optional
-        The axes along which to sum. Uses all axes by default.
-    keepdims : bool, optional
-        Whether or not to keep the dimensions of the original array.
-    dtype: numpy.dtype
-        The data type of the output array.
-
-    Returns
-    -------
-    COO
-        The reduced output sparse array.
-
-    See Also
-    --------
-    :obj:`COO.sum` : Function without ``NaN`` skipping.
-    numpy.nansum : Equivalent Numpy function.
-    """
-    assert out is None
-    x = asCOO(x, name='nansum')
-    return x.nanreduce(np.add, axis=axis, keepdims=keepdims, dtype=dtype)
-
-
-def nanmax(x, axis=None, keepdims=False, dtype=None, out=None):
-    """
-    Maximize along the given axes, skipping ``NaN`` values. Uses all axes by default.
-
-    Parameters
-    ----------
-    x : SparseArray
-        The array to perform the reduction on.
-    axis : Union[int, Iterable[int]], optional
-        The axes along which to maximize. Uses all axes by default.
-    keepdims : bool, optional
-        Whether or not to keep the dimensions of the original array.
-    dtype: numpy.dtype
-        The data type of the output array.
-
-    Returns
-    -------
-    COO
-        The reduced output sparse array.
-
-    See Also
-    --------
-    :obj:`COO.max` : Function without ``NaN`` skipping.
-    numpy.nanmax : Equivalent Numpy function.
-    """
-    assert out is None
-    x = asCOO(x, name='nanmax')
-
-    ar = x.reduce(np.fmax, axis=axis, keepdims=keepdims,
-                  dtype=dtype)
-
-    if (isscalar(ar) and np.isnan(ar)) or np.isnan(ar.data).any():
-        warnings.warn("All-NaN slice encountered", RuntimeWarning, stacklevel=2)
-
-    return ar
-
-
-def nanmin(x, axis=None, keepdims=False, dtype=None, out=None):
-    """
-    Minimize along the given axes, skipping ``NaN`` values. Uses all axes by default.
-
-    Parameters
-    ----------
-    x : SparseArray
-        The array to perform the reduction on.
-    axis : Union[int, Iterable[int]], optional
-        The axes along which to minimize. Uses all axes by default.
-    keepdims : bool, optional
-        Whether or not to keep the dimensions of the original array.
-    dtype: numpy.dtype
-        The data type of the output array.
-
-    Returns
-    -------
-    COO
-        The reduced output sparse array.
-
-    See Also
-    --------
-    :obj:`COO.min` : Function without ``NaN`` skipping.
-    numpy.nanmin : Equivalent Numpy function.
-    """
-    assert out is None
-    x = asCOO(x, name='nanmin')
-
-    ar = x.reduce(np.fmin, axis=axis, keepdims=keepdims,
-                  dtype=dtype)
-
-    if (isscalar(ar) and np.isnan(ar)) or np.isnan(ar.data).any():
-        warnings.warn("All-NaN slice encountered", RuntimeWarning, stacklevel=2)
-
-    return ar
-
-
-def nanprod(x, axis=None, keepdims=False, dtype=None, out=None):
-    """
-    Performs a product operation along the given axes, skipping ``NaN`` values.
-    Uses all axes by default.
-
-    Parameters
-    ----------
-    x : SparseArray
-        The array to perform the reduction on.
-    axis : Union[int, Iterable[int]], optional
-        The axes along which to multiply. Uses all axes by default.
-    keepdims : bool, optional
-        Whether or not to keep the dimensions of the original array.
-    dtype: numpy.dtype
-        The data type of the output array.
-
-    Returns
-    -------
-    COO
-        The reduced output sparse array.
-
-    See Also
-    --------
-    :obj:`COO.prod` : Function without ``NaN`` skipping.
-    numpy.nanprod : Equivalent Numpy function.
-    """
-    assert out is None
-    x = asCOO(x)
-    return x.nanreduce(np.multiply, axis=axis, keepdims=keepdims, dtype=dtype)
diff --git a/sparse/coo/umath.py b/sparse/coo/umath.py
new file mode 100644
index 0000000..f130133
--- /dev/null
+++ b/sparse/coo/umath.py
@@ -0,0 +1,638 @@
+from itertools import product
+
+import numpy as np
+import scipy.sparse
+
+from ..utils import isscalar, PositinalArgumentPartial, _zero_of_dtype
+from ..compatibility import range, zip, zip_longest
+
+
+def elemwise(func, *args, **kwargs):
+    """
+    Apply a function to any number of arguments.
+
+    Parameters
+    ----------
+    func : Callable
+        The function to apply. Must support broadcasting.
+    args : tuple, optional
+        The arguments to the function. Can be :obj:`SparseArray` objects
+        or :obj:`scipy.sparse.spmatrix` objects.
+    kwargs : dict, optional
+        Any additional arguments to pass to the function.
+
+    Returns
+    -------
+    COO
+        The result of applying the function.
+
+    Raises
+    ------
+    ValueError
+        If the operation would result in a dense matrix, or if the operands
+        don't have broadcastable shapes.
+
+    See Also
+    --------
+    :obj:`numpy.ufunc` : A similar Numpy construct. Note that any :code:`ufunc` can be used
+        as the :code:`func` input to this function.
+
+    Notes
+    -----
+    Previously, operations with Numpy arrays were sometimes supported. Now,
+    it is necessary to convert Numpy arrays to :obj:`COO` objects.
+    """
+    # Because we need to mutate args.
+    from .core import COO
+    from ..sparse_array import SparseArray
+
+    args = list(args)
+    posargs = []
+    pos = []
+    for i, arg in enumerate(args):
+        if isinstance(arg, scipy.sparse.spmatrix):
+            args[i] = COO.from_scipy_sparse(arg)
+        elif isscalar(arg) or (isinstance(arg, np.ndarray)
+                               and not arg.shape):
+            # Faster and more reliable to pass ()-shaped ndarrays as scalars.
+            if isinstance(arg, np.ndarray):
+                args[i] = arg[()]
+
+            pos.append(i)
+            posargs.append(args[i])
+        elif isinstance(arg, SparseArray) and not isinstance(arg, COO):
+            args[i] = COO(arg)
+        elif not isinstance(arg, COO):
+            raise ValueError("Performing this operation would produce "
+                             "a dense result: %s" % str(func))
+
+    # Filter out scalars as they are 'baked' into the function.
+    func = PositinalArgumentPartial(func, pos, posargs)
+    args = [arg for arg in args if not isscalar(arg)]
+
+    if len(args) == 0:
+        return func(**kwargs)
+
+    if len(args) == 1:
+        return _elemwise_unary(func, args[0], **kwargs)
+
+    return _elemwise_n_ary(func, *args, **kwargs)
+
+
+# (c) Paul Panzer
+# Taken from https://stackoverflow.com/a/47833496/774273
+# License: https://creativecommons.org/licenses/by-sa/3.0/
+def _match_arrays(a, b):
+    """
+    Finds all indexes into a and b such that a[i] = b[j]. The outputs are sorted
+    in lexographical order.
+
+    Parameters
+    ----------
+    a, b : np.ndarray
+        The input 1-D arrays to match. If matching of multiple fields is
+        needed, use np.recarrays. These two arrays must be sorted.
+
+    Returns
+    -------
+    a_idx, b_idx : np.ndarray
+        The output indices of every possible pair of matching elements.
+    """
+    if len(a) == 0 or len(b) == 0:
+        return np.array([], dtype=np.uint8), np.array([], dtype=np.uint8)
+    asw = np.r_[0, 1 + np.flatnonzero(a[:-1] != a[1:]), len(a)]
+    bsw = np.r_[0, 1 + np.flatnonzero(b[:-1] != b[1:]), len(b)]
+    al, bl = np.diff(asw), np.diff(bsw)
+    na = len(al)
+    asw, bsw = asw, bsw
+    abunq = np.r_[a[asw[:-1]], b[bsw[:-1]]]
+    m = np.argsort(abunq, kind='mergesort')
+    mv = abunq[m]
+    midx = np.flatnonzero(mv[:-1] == mv[1:])
+    ai, bi = m[midx], m[midx + 1] - na
+    aic = np.r_[0, np.cumsum(al[ai])]
+    a_idx = np.ones((aic[-1],), dtype=np.int_)
+    a_idx[aic[:-1]] = asw[ai]
+    a_idx[aic[1:-1]] -= asw[ai[:-1]] + al[ai[:-1]] - 1
+    a_idx = np.repeat(np.cumsum(a_idx), np.repeat(bl[bi], al[ai]))
+    bi = np.repeat(bi, al[ai])
+    bic = np.r_[0, np.cumsum(bl[bi])]
+    b_idx = np.ones((bic[-1],), dtype=np.int_)
+    b_idx[bic[:-1]] = bsw[bi]
+    b_idx[bic[1:-1]] -= bsw[bi[:-1]] + bl[bi[:-1]] - 1
+    b_idx = np.cumsum(b_idx)
+    return a_idx, b_idx
+
+
+def _elemwise_n_ary(func, *args, **kwargs):
+    """
+    Apply a function to any number of arguments with broadcasting.
+
+    Parameters
+    ----------
+    func : Callable
+        The function to apply to arguments. Must support broadcasting.
+    args : list
+        Input :obj:`COO` or :obj:`numpy.ndarray`s.
+    kwargs : dict
+        Additional arguments to pass to the function.
+
+    Returns
+    -------
+    COO
+        The output array.
+
+    Raises
+    ------
+    ValueError
+        If the input shapes aren't compatible or the result will be dense.
+    """
+    from .core import COO
+
+    args = list(args)
+
+    for arg in args:
+        if isinstance(arg, COO):
+            arg.sum_duplicates()
+
+    args_zeros = tuple(_zero_of_dtype(arg.dtype)[()] for arg in args)
+
+    func_value = func(*args_zeros, **kwargs)
+    func_zero = _zero_of_dtype(func_value.dtype)
+    if func_value != func_zero:
+        raise ValueError("Performing this operation would produce "
+                         "a dense result: %s" % str(func))
+    data_list = []
+    coords_list = []
+
+    cache = {}
+    for mask in product([True, False], repeat=len(args)):
+        if not any(mask):
+            continue
+
+        ci, di = _unmatch_coo(func, args, mask, cache, **kwargs)
+
+        coords_list.extend(ci)
+        data_list.extend(di)
+
+    result_shape = _get_nary_broadcast_shape(*[arg.shape for arg in args])
+
+    # Concatenate matches and mismatches
+    data = np.concatenate(data_list) if len(data_list) else np.empty((0,), dtype=args[0].dtype)
+    coords = np.concatenate(coords_list, axis=1) if len(coords_list) else \
+        np.empty((0, len(result_shape)), dtype=np.min_scalar_type(max(result_shape) - 1))
+
+    nonzero = data != func_zero
+    data = data[nonzero]
+    coords = coords[:, nonzero]
+
+    return COO(coords, data, shape=result_shape, has_duplicates=False)
+
+
+def _match_coo(*args, **kwargs):
+    """
+    Matches the coordinates for any number of input :obj:`COO` arrays.
+    Equivalent to "sparse" broadcasting for all arrays.
+
+    Parameters
+    ----------
+    args : Tuple[COO]
+        The input :obj:`COO` arrays.
+    return_midx : bool
+        Whether to return matched indices or matched arrays. Matching
+        only supported for two arrays. ``False`` by default.
+    cache : dict
+        Cache of things already matched. No cache by default.
+
+    Returns
+    -------
+    matched_idx : List[ndarray]
+        The indices of matched elements in the original arrays. Only returned if
+        ``return_midx`` is ``True``.
+    matched_arrays : List[COO]
+        The expanded, matched :obj:`COO` objects. Only returned if
+        ``return_midx`` is ``False``.
+    """
+    from .core import COO
+    from .common import linear_loc
+
+    return_midx = kwargs.pop('return_midx', False)
+    cache = kwargs.pop('cache', None)
+
+    if kwargs:
+        raise ValueError('Unknown kwargs %s' % kwargs.keys())
+
+    if return_midx and (len(args) != 2 or cache is not None):
+        raise NotImplementedError('Matching indices only supported for two args, and no cache.')
+
+    matched_arrays = [args[0]]
+    cache_key = [id(args[0])]
+    for arg2 in args[1:]:
+        cache_key.append(id(arg2))
+        key = tuple(cache_key)
+        if cache is not None and key in cache:
+            matched_arrays = cache[key]
+            continue
+
+        cargs = [matched_arrays[0], arg2]
+        current_shape = _get_broadcast_shape(matched_arrays[0].shape, arg2.shape)
+        params = [_get_broadcast_parameters(arg.shape, current_shape) for arg in cargs]
+        reduced_params = [all(p) for p in zip(*params)]
+        reduced_shape = _get_reduced_shape(arg2.shape,
+                                           reduced_params[-arg2.ndim:])
+
+        reduced_coords = [_get_reduced_coords(arg.coords, reduced_params[-arg.ndim:])
+                          for arg in cargs]
+
+        linear = [linear_loc(rc, reduced_shape) for rc in reduced_coords]
+        sorted_idx = [np.argsort(idx) for idx in linear]
+        linear = [idx[s] for idx, s in zip(linear, sorted_idx)]
+        matched_idx = _match_arrays(*linear)
+
+        if return_midx:
+            matched_idx = [sidx[midx] for sidx, midx in zip(sorted_idx, matched_idx)]
+            return matched_idx
+
+        coords = [arg.coords[:, s] for arg, s in zip(cargs, sorted_idx)]
+        mcoords = [c[:, idx] for c, idx in zip(coords, matched_idx)]
+        mcoords = _get_matching_coords(mcoords, params, current_shape)
+        mdata = [arg.data[sorted_idx[0]][matched_idx[0]] for arg in matched_arrays]
+        mdata.append(arg2.data[sorted_idx[1]][matched_idx[1]])
+        matched_arrays = [COO(mcoords, md, shape=current_shape) for md in mdata]
+
+        if cache is not None:
+            cache[key] = matched_arrays
+
+    return matched_arrays
+
+
+def _unmatch_coo(func, args, mask, cache, **kwargs):
+    """
+    Matches the coordinates for any number of input :obj:`COO` arrays.
+
+    First computes the matches, then filters out the non-matches.
+
+    Parameters
+    ----------
+    func : Callable
+        The function to compute matches
+    args : tuple[COO]
+        The input :obj:`COO` arrays.
+    mask : tuple[bool]
+        Specifies the inputs that are zero and the ones that are
+        nonzero.
+    kwargs: dict
+        Extra keyword arguments to pass to func.
+
+    Returns
+    -------
+    matched_coords : list[ndarray]
+        The matched coordinates.
+    matched_data : list[ndarray]
+        The matched data.
+    """
+    from .core import COO
+
+    matched_args = [a for a, m in zip(args, mask) if m]
+    unmatched_args = [a for a, m in zip(args, mask) if not m]
+
+    matched_arrays = _match_coo(*matched_args, cache=cache)
+
+    pos = tuple(i for i, m in enumerate(mask) if not m)
+    posargs = [_zero_of_dtype(arg.dtype)[()] for arg, m in zip(args, mask) if not m]
+    result_shape = _get_nary_broadcast_shape(*[arg.shape for arg in args])
+
+    partial = PositinalArgumentPartial(func, pos, posargs)
+    matched_func = partial(*[a.data for a in matched_arrays], **kwargs)
+
+    unmatched_mask = matched_func != _zero_of_dtype(matched_func.dtype)
+
+    if not unmatched_mask.any():
+        return [], []
+
+    func_data = matched_func[unmatched_mask]
+    func_coords = matched_arrays[0].coords[:, unmatched_mask]
+
+    func_array = COO(func_coords, func_data, shape=matched_arrays[0].shape).broadcast_to(result_shape)
+
+    if all(mask):
+        return [func_array.coords], [func_array.data]
+
+    unmatched_mask = np.ones(func_array.nnz, dtype=np.bool)
+
+    for arg in unmatched_args:
+        matched_idx = _match_coo(func_array, arg, return_midx=True)[0]
+        unmatched_mask[matched_idx] = False
+
+    coords = np.asarray(func_array.coords[:, unmatched_mask], order='C')
+    data = np.asarray(func_array.data[unmatched_mask], order='C')
+
+    return [coords], [data]
+
+
+def _get_nary_broadcast_shape(*shapes):
+    """
+    Broadcast any number of shapes to a result shape.
+
+    Parameters
+    ----------
+    shapes : tuple[tuple[int]]
+        The shapes to broadcast.
+
+    Returns
+    -------
+    tuple[int]
+        The output shape.
+
+    Raises
+    ------
+    ValueError
+        If the input shapes cannot be broadcast to a single shape.
+    """
+    result_shape = ()
+
+    for shape in shapes:
+        try:
+            result_shape = _get_broadcast_shape(shape, result_shape)
+        except ValueError:
+            shapes_str = ', '.join(str(shape) for shape in shapes)
+            raise ValueError('operands could not be broadcast together with shapes %s'
+                             % shapes_str)
+
+    return result_shape
+
+
+def _get_broadcast_shape(shape1, shape2, is_result=False):
+    """
+    Get the overall broadcasted shape.
+
+    Parameters
+    ----------
+    shape1, shape2 : tuple[int]
+        The input shapes to broadcast together.
+    is_result : bool
+        Whether or not shape2 is also the result shape.
+
+    Returns
+    -------
+    result_shape : tuple[int]
+        The overall shape of the result.
+
+    Raises
+    ------
+    ValueError
+        If the two shapes cannot be broadcast together.
+    """
+    # https://stackoverflow.com/a/47244284/774273
+    if not all((l1 == l2) or (l1 == 1) or ((l2 == 1) and not is_result) for l1, l2 in
+               zip(shape1[::-1], shape2[::-1])):
+        raise ValueError('operands could not be broadcast together with shapes %s, %s' %
+                         (shape1, shape2))
+
+    result_shape = tuple(max(l1, l2) for l1, l2 in
+                         zip_longest(shape1[::-1], shape2[::-1], fillvalue=1))[::-1]
+
+    return result_shape
+
+
+def _get_broadcast_parameters(shape, broadcast_shape):
+    """
+    Get the broadcast parameters.
+
+    Parameters
+    ----------
+    shape : tuple[int]
+        The input shape.
+    broadcast_shape
+        The shape to broadcast to.
+
+    Returns
+    -------
+    params : list
+        A list containing None if the dimension isn't in the original array, False if
+        it needs to be broadcast, and True if it doesn't.
+    """
+    params = [None if l1 is None else l1 == l2 for l1, l2
+              in zip_longest(shape[::-1], broadcast_shape[::-1], fillvalue=None)][::-1]
+
+    return params
+
+
+def _get_reduced_coords(coords, params):
+    """
+    Gets only those dimensions of the coordinates that don't need to be broadcast.
+
+    Parameters
+    ----------
+    coords : np.ndarray
+        The coordinates to reduce.
+    params : list
+        The params from which to check which dimensions to get.
+
+    Returns
+    -------
+    reduced_coords : np.ndarray
+        The reduced coordinates.
+    """
+
+    reduced_params = [bool(param) for param in params]
+
+    return coords[reduced_params]
+
+
+def _get_reduced_shape(shape, params):
+    """
+    Gets only those dimensions of the coordinates that don't need to be broadcast.
+
+    Parameters
+    ----------
+    coords : np.ndarray
+        The coordinates to reduce.
+    params : list
+        The params from which to check which dimensions to get.
+
+    Returns
+    -------
+    reduced_coords : np.ndarray
+        The reduced coordinates.
+    """
+    reduced_shape = tuple(l for l, p in zip(shape, params) if p)
+
+    return reduced_shape
+
+
+def _get_expanded_coords_data(coords, data, params, broadcast_shape):
+    """
+    Expand coordinates/data to broadcast_shape. Does most of the heavy lifting for broadcast_to.
+    Produces sorted output for sorted inputs.
+
+    Parameters
+    ----------
+    coords : np.ndarray
+        The coordinates to expand.
+    data : np.ndarray
+        The data corresponding to the coordinates.
+    params : list
+        The broadcast parameters.
+    broadcast_shape : tuple[int]
+        The shape to broadcast to.
+
+    Returns
+    -------
+    expanded_coords : np.ndarray
+        List of 1-D arrays. Each item in the list has one dimension of coordinates.
+    expanded_data : np.ndarray
+        The data corresponding to expanded_coords.
+    """
+    first_dim = -1
+    expand_shapes = []
+    for d, p, l in zip(range(len(broadcast_shape)), params, broadcast_shape):
+        if p and first_dim == -1:
+            expand_shapes.append(coords.shape[1])
+            first_dim = d
+
+        if not p:
+            expand_shapes.append(l)
+
+    all_idx = _cartesian_product(*(np.arange(d, dtype=np.min_scalar_type(d - 1)) for d in expand_shapes))
+    dt = np.result_type(*(np.min_scalar_type(l - 1) for l in broadcast_shape))
+
+    false_dim = 0
+    dim = 0
+
+    expanded_coords = np.empty((len(broadcast_shape), all_idx.shape[1]), dtype=dt)
+    expanded_data = data[all_idx[first_dim]]
+
+    for d, p, l in zip(range(len(broadcast_shape)), params, broadcast_shape):
+        if p:
+            expanded_coords[d] = coords[dim, all_idx[first_dim]]
+        else:
+            expanded_coords[d] = all_idx[false_dim + (d > first_dim)]
+            false_dim += 1
+
+        if p is not None:
+            dim += 1
+
+    return np.asarray(expanded_coords), np.asarray(expanded_data)
+
+
+# (c) senderle
+# Taken from https://stackoverflow.com/a/11146645/774273
+# License: https://creativecommons.org/licenses/by-sa/3.0/
+def _cartesian_product(*arrays):
+    """
+    Get the cartesian product of a number of arrays.
+
+    Parameters
+    ----------
+    arrays : Tuple[np.ndarray]
+        The arrays to get a cartesian product of. Always sorted with respect
+        to the original array.
+
+    Returns
+    -------
+    out : np.ndarray
+        The overall cartesian product of all the input arrays.
+    """
+    broadcastable = np.ix_(*arrays)
+    broadcasted = np.broadcast_arrays(*broadcastable)
+    rows, cols = np.prod(broadcasted[0].shape), len(broadcasted)
+    dtype = np.result_type(*arrays)
+    out = np.empty(rows * cols, dtype=dtype)
+    start, end = 0, rows
+    for a in broadcasted:
+        out[start:end] = a.reshape(-1)
+        start, end = end, end + rows
+    return out.reshape(cols, rows)
+
+
+def _elemwise_unary(func, self, *args, **kwargs):
+    from .core import COO
+
+    check = kwargs.pop('check', True)
+    data_zero = _zero_of_dtype(self.dtype)
+    zero_func = func(data_zero, *args, **kwargs)
+    func_zero = _zero_of_dtype(zero_func.dtype)
+    if check and zero_func != func_zero:
+        raise ValueError("Performing this operation would produce "
+                         "a dense result: %s" % str(func))
+
+    data_func = func(self.data, *args, **kwargs)
+    nonzero = data_func != func_zero
+
+    return COO(self.coords[:, nonzero], data_func[nonzero],
+               shape=self.shape,
+               has_duplicates=self.has_duplicates,
+               sorted=self.sorted)
+
+
+def _get_matching_coords(coords, params, shape):
+    """
+    Get the matching coords across a number of broadcast operands.
+
+    Parameters
+    ----------
+    coords : list[numpy.ndarray]
+        The input coordinates.
+    params : list[Union[bool, none]]
+        The broadcast parameters.
+    Returns
+    -------
+    numpy.ndarray
+        The broacasted coordinates
+    """
+    matching_coords = []
+    dims = np.zeros(len(coords), dtype=np.uint8)
+
+    for p_all in zip(*params):
+        for i, p in enumerate(p_all):
+            if p:
+                matching_coords.append(coords[i][dims[i]])
+                break
+        else:
+            matching_coords.append(coords[dims[0]])
+
+        for i, p in enumerate(p_all):
+            if p is not None:
+                dims[i] += 1
+
+    dtype = np.min_scalar_type(max(shape) - 1)
+
+    return np.asarray(matching_coords, dtype=dtype)
+
+
+def broadcast_to(x, shape):
+    """
+    Performs the equivalent of :obj:`numpy.broadcast_to` for :obj:`COO`. Note that
+    this function returns a new array instead of a view.
+
+    Parameters
+    ----------
+    shape : tuple[int]
+        The shape to broadcast the data to.
+
+    Returns
+    -------
+    COO
+        The broadcasted sparse array.
+
+    Raises
+    ------
+    ValueError
+        If the operand cannot be broadcast to the given shape.
+
+    See also
+    --------
+    :obj:`numpy.broadcast_to` : NumPy equivalent function
+    """
+    from .core import COO
+
+    if shape == x.shape:
+        return x
+
+    result_shape = _get_broadcast_shape(x.shape, shape, is_result=True)
+    params = _get_broadcast_parameters(x.shape, result_shape)
+    coords, data = _get_expanded_coords_data(x.coords, x.data, params, result_shape)
+
+    return COO(coords, data, shape=result_shape, has_duplicates=x.has_duplicates,
+               sorted=x.sorted)


Test patch:
diff --git a/sparse/tests/test_coo.py b/sparse/tests/test_coo.py
index 8c4a767..d903683 100644
--- a/sparse/tests/test_coo.py
+++ b/sparse/tests/test_coo.py
@@ -441,7 +441,7 @@ def test_trinary_broadcasting_pathological(shapes, func, value, fraction):
 
 
 def test_sparse_broadcasting(monkeypatch):
-    orig_unmatch_coo = sparse.coo._unmatch_coo
+    orig_unmatch_coo = sparse.coo.umath._unmatch_coo
 
     state = {'num_matches': 0}
 
@@ -453,7 +453,7 @@ def test_sparse_broadcasting(monkeypatch):
         state['num_matches'] += len(result[0])
         return result
 
-    monkeypatch.setattr(sparse.coo, '_unmatch_coo', mock_unmatch_coo)
+    monkeypatch.setattr(sparse.coo.umath, '_unmatch_coo', mock_unmatch_coo)
 
     xs * ys
 
@@ -462,7 +462,7 @@ def test_sparse_broadcasting(monkeypatch):
 
 
 def test_dense_broadcasting(monkeypatch):
-    orig_unmatch_coo = sparse.coo._unmatch_coo
+    orig_unmatch_coo = sparse.coo.umath._unmatch_coo
 
     state = {'num_matches': 0}
 
@@ -474,7 +474,7 @@ def test_dense_broadcasting(monkeypatch):
         state['num_matches'] += len(result[0])
         return result
 
-    monkeypatch.setattr(sparse.coo, '_unmatch_coo', mock_unmatch_coo)
+    monkeypatch.setattr(sparse.coo.umath, '_unmatch_coo', mock_unmatch_coo)
 
     xs + ys