Repository: yt-project/unyt
Pull request #200 (2021-08-19 16:39:43)

ENH: Provisional support for NEP 18 (__array_function__ protocol)
My initial motivation here was to add some unit representation to error messages when comparing two unyt_array instance via functions from `numpy.testing`, like so:

```python
import numpy as np
import unyt as un

a = [1, 2, 3] * un.cm
b = [1, 2, 3] * un.km
np.testing.assert_array_equal(a, b)
```
which yields, on master:

```
...
AssertionError:
Arrays are not equal

Mismatched elements: 3 / 3 (100%)
Max absolute difference: 299997.
Max relative difference: 0.99999
 x: unyt_array([1, 2, 3])
 y: unyt_array([1, 2, 3])
```
and on this branch
<details>
<summary>previous version</summary>

```
...
AssertionError:
Arrays are not equal

Mismatched elements: 3 / 3 (100%)
Max absolute difference: 299997. cm
Max relative difference: 0.99999 dimensionless
 x: unyt_array([1, 2, 3] cm)
 y: unyt_array([1, 2, 3] km)
```

</details>

edit:
```
AssertionError:
Arrays are not equal

Mismatched elements: 3 / 3 (100%)
Max absolute difference: 299997., units='cm'
Max relative difference: 0.99999, units='dimensionless'
 x: unyt_array([1, 2, 3], units='cm')
 y: unyt_array([1, 2, 3], units='km')
```

Incidentally, it turns out that fixing this necessitated a kick-off implementation of NEP 18, so this work laid the foundation to solve:
- [x] #69
- [x] #130
- [ ] #50 (most likely out of scope)

More broadly, implementing NEP 18 for any was the topic of #139.
Granted I need to take more time to check that I'm not going against the original intentions here. My current approach is that, since covering the whole numpy public API in one go seems like an gigantic task, I'm implementing `unyt_array.__array_function__` with a fallthrough condition: if a special case isn't implemented yet, just fallback to the raw numpy implem (which is currently the behaviour for _all_ functions subject to NEP 18). This way we can add support for more and more functions in a progressive way.
I'm going to set the bar low(ish) for now, and try and fix the already reported cases, as reported above, as a first step.


An important question to address is: what should be done in the general case where we don't have a custom implementation for an array function ?

1) transparently default to the raw numpy implementation without a warning (this is de facto what is done as of unyt 2.8.0, and will remain the case until NEP 18 isn't at least partially implemented).
2) same as 1. _but_ emit a warning (possibly, we could have a whitelist of functions that are known to be perfectly fine without a specialised implementation, for which no warning would be emitted) along the lines of 
```
UserWarning: calling `numpy.FUNC` on a unyt_array. Results may hold incorrect units. A future version of unyt will remove this warning, and possibly change the behaviour of this function to be dimensionally correct.
```
3) Error out

Option 1 this is the current implementation of this PR because I think it is the less disruptive or noisy one. My personal opinion is that it's probably okay to have incomplete support for NEP 18 for a couple release, as long as it is clearly stated in the release notes.

neutrinoceros: I've gone the long way and started implemented a bunch of functions with obvious requirements.
TODO:

## inverse-like functions
- [x] `np.linalg.inv`
- [x] `np.linalg.tensorinv`
- [ ] `np.linalg.pinv` (done-ish, more testing needed for the rcond parameter)

## product-like functions
- [x] `np.dot`
- [ ] `np.inner`
- [ ] `np.outer`
- [ ] `np.matmul`

## histogram functions (as requested for #130)
- [x] `np.histogram`
- [x] `np.histogram2d`
- [x] `np.histogramdd`

neutrinoceros: pre-commit.ci autofix
neutrinoceros: @jzuhone do you think you'd be able to review this ? Is there someone else I could ask ?
jzuhone: @neutrinoceros I can try to look at it next week
neutrinoceros: @jzuhone I would like to propose we bundle this PR in the next release alongside the dask layer. IMO, each of these would justify a release on their own, so it's up to you :)
jzuhone: @chrishavlin or @matthewturk could you look at this one?
ngoldbaum: I'm a little concerned that the coverage of numpy functions isn't as extensive as users might like and that implementing support in a piecemeal fashion might lead to more breakage than users would prefer, since fundamentally we're changing the type of the outputs of numpy functions. If we're going to break that, it would be best to break them all at once instead of staggering that over many releases as support gets added in a piecemeal fashion.

unyt ships with a number of wrappers for numpy functions (`uconcatenate`, `ucross`, `udot`, etc). We should make sure we have replacements for all of those (since there's clearly already a need for those) and deprecate the wrappers. I know at least yt uses those wrappers.

If you want a hand with this I'm happy to help out.
neutrinoceros: I agree that it'd be best to do it all in one clean go, and yes I would love a hand here. How should we proceed ? If you want to keep it all in a single PR feel free to push to my branch.
ngoldbaum: Sure I’ll push to this branch
neutrinoceros: @ngoldbaum do you know of a reliable way to generate the complete list of functions we should be covering ?
ngoldbaum: Yes, see https://github.com/numpy/numpy/issues/15544#issuecomment-584453408 as well as other comments in that issue
ngoldbaum: Can you make sure that the existing wrappers in unyt that this PR obsoletes (e.g. `udot`) are deprecated? You could probably rewrite them to be thin wrappers around the regular numpy functions now.
neutrinoceros: good call, I completely missed that. I think it may be preferable _not_ to rewrite them at all in case there's _any_ unanticipated difference so users have a chance to report and still get something that works for them while we fix the replacement, but I'll definitely add those deprecation warnings.
neutrinoceros: I rebased the branch but didn't write the history. I added two commit, the first of which corresponds to #288, and the second one adds deprecation warnings as requested.
 
neutrinoceros: pre-commit.ci autofix

Patch:
diff --git a/pyproject.toml b/pyproject.toml
index 683266c..50f5143 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -31,9 +31,3 @@ skip =  ["venv", "benchmarks"]
 [tool.setuptools_scm]
 write_to = "unyt/_version.py"
 version_scheme = "post-release"
-
-[tool.pytest.ini_options]
-addopts = "--ignore=benchmarks --ignore=paper --ignore=unyt/_mpl_array_converter"
-filterwarnings = [
-    "error",
-]
diff --git a/unyt/_array_functions.py b/unyt/_array_functions.py
new file mode 100644
index 0000000..f3e9899
--- /dev/null
+++ b/unyt/_array_functions.py
@@ -0,0 +1,188 @@
+import numpy as np
+
+from unyt.exceptions import UnitConversionError
+
+_HANDLED_FUNCTIONS = {}
+
+
+def implements(numpy_function):
+    """Register an __array_function__ implementation for unyt_array objects."""
+    # See NEP 18 https://numpy.org/neps/nep-0018-array-function-protocol.html
+    def decorator(func):
+        _HANDLED_FUNCTIONS[numpy_function] = func
+        return func
+
+    return decorator
+
+
+@implements(np.array2string)
+def array2string(a, *args, **kwargs):
+    return (
+        np.array2string._implementation(a, *args, **kwargs)
+        + f", units={str(a.units)!r}"
+    )
+
+
+def _get_conversion_factor(out, units) -> float:
+    # helper function to "product"-like functions with an 'out' argument
+    try:
+        return (1 * out.units).to(units).d
+    except (AttributeError, UnitConversionError) as exc:
+        raise TypeError(
+            "output array is not acceptable "
+            f"(units '{out.units}' cannot be converted to '{units}')"
+        ) from exc
+
+
+@implements(np.dot)
+def dot(a, b, out=None):
+    prod_units = a.units * b.units
+    if out is None:
+        return np.dot._implementation(a.ndview, b.ndview) * prod_units
+
+    conv_factor = _get_conversion_factor(out, prod_units)
+
+    np.dot._implementation(a.ndview, b.ndview, out=out.ndview)
+    if not out.units == prod_units:
+        out[:] *= conv_factor
+    return out
+
+
+@implements(np.vdot)
+def vdot(a, b):
+    return np.vdot._implementation(a.ndview, b.ndview) * (a.units * b.units)
+
+
+@implements(np.inner)
+def inner(a, b):
+    return np.inner._implementation(a.ndview, b.ndview) * (a.units * b.units)
+
+
+@implements(np.outer)
+def outer(a, b, out=None):
+    prod_units = a.units * b.units
+    if out is None:
+        return np.outer._implementation(a.ndview, b.ndview) * prod_units
+
+    conv_factor = _get_conversion_factor(out, prod_units)
+
+    np.outer._implementation(a.ndview, b.ndview, out=out.ndview)
+    if not out.units == prod_units:
+        out[:] *= conv_factor
+    return out
+
+
+@implements(np.kron)
+def kron(a, b):
+    return np.kron._implementation(a.ndview, b.ndview) * (a.units * b.units)
+
+
+@implements(np.matmul)
+def matmul(a, b, out=None, **kwargs):
+    prod_units = a.units * b.units
+    if out is None:
+        return np.matmul._implementation(a.ndview, b.ndview) * prod_units
+
+    conv_factor = _get_conversion_factor(out, prod_units)
+    np.matmul._implementation(a.ndview, b.ndview, out=out.ndview)
+    if not out.units == prod_units:
+        out[:] *= conv_factor
+    return out
+
+
+@implements(np.linalg.inv)
+def linalg_inv(a, *args, **kwargs):
+    return np.linalg.inv._implementation(a.ndview, *args, **kwargs) / a.units
+
+
+@implements(np.linalg.tensorinv)
+def linalg_tensorinv(a, *args, **kwargs):
+    return np.linalg.tensorinv._implementation(a, *args, **kwargs) / a.units
+
+
+@implements(np.linalg.pinv)
+def linalg_pinv(a, *args, **kwargs):
+    return np.linalg.pinv._implementation(a, *args, **kwargs).ndview / a.units
+
+
+def _sanitize_range(_range, units):
+    # helper function to histogram* functions
+    ndim = len(units)
+    if _range is None:
+        return _range
+    new_range = np.empty((ndim, 2))
+    for i in range(ndim):
+        ilim = _range[2 * i : 2 * (i + 1)]
+        imin, imax = ilim
+        if not (hasattr(imin, "units") and hasattr(imax, "units")):
+            raise TypeError(
+                f"Elements of range must both have a 'units' attribute. Got {_range}"
+            )
+        new_range[i] = imin.to(units[i]).value, imax.to(units[i]).value
+    return new_range.squeeze()
+
+
+@implements(np.histogram)
+def histogram(
+    a,
+    bins=10,
+    range=None,
+    normed=None,
+    weights=None,
+    density=None,
+):
+    range = _sanitize_range(range, units=[a.units])
+    counts, bins = np.histogram._implementation(
+        a.ndview,
+        bins=bins,
+        range=range,
+        normed=normed,
+        weights=weights,
+        density=density,
+    )
+    return counts, bins * a.units
+
+
+@implements(np.histogram2d)
+def histogram2d(
+    x,
+    y,
+    bins=10,
+    range=None,
+    normed=None,
+    weights=None,
+    density=None,
+):
+    range = _sanitize_range(range, units=[x.units, y.units])
+    counts, xbins, ybins = np.histogram2d._implementation(
+        x.ndview,
+        y.ndview,
+        bins=bins,
+        range=range,
+        normed=normed,
+        weights=weights,
+        density=density,
+    )
+    return counts, xbins * x.units, ybins * y.units
+
+
+@implements(np.histogramdd)
+def histogramdd(
+    sample,
+    bins=10,
+    range=None,
+    normed=None,
+    weights=None,
+    density=None,
+):
+    units = [_.units for _ in sample]
+    range = _sanitize_range(range, units=units)
+    counts, bins = np.histogramdd._implementation(
+        [_.ndview for _ in sample],
+        bins=bins,
+        range=range,
+        normed=normed,
+        weights=weights,
+        density=density,
+    )
+    return counts, tuple(_bin * u for _bin, u in zip(bins, units))
diff --git a/unyt/_deprecation.py b/unyt/_deprecation.py
new file mode 100644
index 0000000..4660eec
--- /dev/null
+++ b/unyt/_deprecation.py
@@ -0,0 +1,22 @@
+import warnings
+from typing import Optional
+
+
+def warn_deprecated(
+    name,
+    /,
+    *,
+    stacklevel: int = 3,
+    replacement: Optional[str] = None,
+    since_version: str,
+) -> None:
+    msg = (
+        f"{name} is deprecated and will be removed in a future version\n"
+        f"Instead, {replacement}\n"
+        f"(deprecated since unyt v{since_version})"
+    )
+    warnings.warn(
+        msg,
+        category=DeprecationWarning,
+        stacklevel=stacklevel,
+    )
diff --git a/unyt/array.py b/unyt/array.py
index de2b534..02ab8f5 100644
--- a/unyt/array.py
+++ b/unyt/array.py
@@ -110,6 +110,7 @@ from numpy import (
 from numpy.core.umath import _ones_like, clip
 from sympy import Rational
 
+from unyt._array_functions import _HANDLED_FUNCTIONS
 from unyt._on_demand_imports import _astropy, _dask, _pint
 from unyt._pint_conversions import convert_pint_units
 from unyt._unit_lookup_table import default_unit_symbol_lut
@@ -133,6 +134,8 @@ from unyt.unit_registry import (
     default_unit_registry,
 )
 
+from ._deprecation import warn_deprecated
+
 NULL_UNIT = Unit()
 POWER_MAPPING = {multiply: lambda x: x, divide: lambda x: 2 - x}
 
@@ -1947,6 +1950,18 @@ class unyt_array(np.ndarray):
             return out_arr
         return mul * out_arr
 
+    def __array_function__(self, func, types, args, kwargs):
+        # Follow NEP 18 guidelines
+        # https://numpy.org/neps/nep-0018-array-function-protocol.html
+        if func not in _HANDLED_FUNCTIONS:
+            # default to numpy's private implementation
+            return func._implementation(*args, **kwargs)
+        # Note: this allows subclasses that don't override
+        # __array_function__ to handle unyt_array objects
+        if not all(issubclass(t, unyt_array) for t in types):
+            return NotImplemented
+        return _HANDLED_FUNCTIONS[func](*args, **kwargs)
+
     def copy(self, order="C"):
         """
         Return a copy of the array.
@@ -2203,6 +2218,10 @@ def uconcatenate(arrs, axis=0):
     unyt_array([1, 2, 3, 2, 3, 4], 'cm')
 
     """
+    # TODO: see https://github.com/yt-project/unyt/issues/289
+    # warn_deprecated(
+    #    "unyt.uconcatenate", replacement="use numpy.concatenate", since_version="2.10"
+    # )
     v = np.concatenate(arrs, axis=axis)
     v = _validate_numpy_wrapper_units(v, arrs)
     return v
@@ -2215,6 +2234,10 @@ def ucross(arr1, arr2, registry=None, axisa=-1, axisb=-1, axisc=-1, axis=None):
     See the documentation of numpy.cross for full
     details.
     """
+    # TODO: see https://github.com/yt-project/unyt/issues/289
+    # warn_deprecated(
+    #    "unyt.ucross", replacement="use numpy.cross", since_version="2.10"
+    # )
     v = np.cross(arr1, arr2, axisa=axisa, axisb=axisb, axisc=axisc, axis=axis)
     units = arr1.units * arr2.units
     arr = unyt_array(v, units, registry=registry)
@@ -2237,6 +2260,10 @@ def uintersect1d(arr1, arr2, assume_unique=False):
     unyt_array([2, 3], 'cm')
 
     """
+    # TODO: see https://github.com/yt-project/unyt/issues/289
+    # warn_deprecated(
+    #    "unyt.uintersect1d", replacement="use numpy.intersect1d", since_version="2.10"
+    # )
     v = np.intersect1d(arr1, arr2, assume_unique=assume_unique)
     v = _validate_numpy_wrapper_units(v, [arr1, arr2])
     return v
@@ -2258,6 +2285,10 @@ def uunion1d(arr1, arr2):
     unyt_array([1, 2, 3, 4], 'cm')
 
     """
+    # TODO: see https://github.com/yt-project/unyt/issues/289
+    # warn_deprecated(
+    #    "unyt.uunion1d", replacement="use numpy.union1d", since_version="2.10"
+    # )
     v = np.union1d(arr1, arr2)
     v = _validate_numpy_wrapper_units(v, [arr1, arr2])
     return v
@@ -2277,6 +2308,10 @@ def unorm(data, ord=None, axis=None, keepdims=False):
     >>> print(unorm(data))
     3.7416573867739413 km
     """
+    # TODO: see https://github.com/yt-project/unyt/issues/289
+    # warn_deprecated(
+    #    "unyt.unorm", replacement="use numpy.norm", since_version="2.10"
+    # )
     norm = np.linalg.norm(data, ord=ord, axis=axis, keepdims=keepdims)
     if norm.shape == ():
         return unyt_quantity(norm, data.units)
@@ -2297,6 +2332,7 @@ def udot(op1, op2):
     [[2. 2.]
      [2. 2.]] km*s
     """
+    warn_deprecated("unyt.udot", replacement="use numpy.dot", since_version="2.10")
     dot = np.dot(op1.d, op2.d)
     units = op1.units * op2.units
     if dot.shape == ():
@@ -2318,6 +2354,10 @@ def uvstack(arrs):
     [[1 2 3]
      [2 3 4]] km
     """
+    # TODO: see https://github.com/yt-project/unyt/issues/289
+    # warn_deprecated(
+    #    "unyt.uvstack", replacement="use numpy.vstack", since_version="2.10"
+    # )
     v = np.vstack(arrs)
     v = _validate_numpy_wrapper_units(v, arrs)
     return v
@@ -2342,6 +2382,10 @@ def uhstack(arrs):
      [2 3]
      [3 4]] km
     """
+    # TODO: see https://github.com/yt-project/unyt/issues/289
+    # warn_deprecated(
+    #    "unyt.uhstack", replacement="use numpy.hstack", since_version="2.10"
+    # )
     v = np.hstack(arrs)
     v = _validate_numpy_wrapper_units(v, arrs)
     return v
@@ -2366,6 +2410,10 @@ def ustack(arrs, axis=0):
     [[1 2 3]
      [2 3 4]] km
     """
+    # TODO: see https://github.com/yt-project/unyt/issues/289
+    # warn_deprecated(
+    #    "unyt.ustack", replacement="use numpy.stack", since_version="2.10"
+    # )
     v = np.stack(arrs, axis=axis)
     v = _validate_numpy_wrapper_units(v, arrs)
     return v
@@ -2568,6 +2616,10 @@ def allclose_units(actual, desired, rtol=1e-7, atol=0, **kwargs):
     >>> allclose_units(actual, desired)
     True
     """
+    # TODO: see https://github.com/yt-project/unyt/issues/289
+    # warn_deprecated(
+    #    "unyt.allclose_units", replacement="use numpy.allclose", since_version="2.10"
+    # )
     # Create a copy to ensure this function does not alter input arrays
     act = unyt_array(actual)
     des = unyt_array(desired)


Test patch:
diff --git a/unyt/tests/test_array_functions.py b/unyt/tests/test_array_functions.py
new file mode 100644
index 0000000..81c5dbf
--- /dev/null
+++ b/unyt/tests/test_array_functions.py
@@ -0,0 +1,201 @@
+# tests for NEP 18
+import re
+
+import numpy as np
+import pytest
+
+from unyt import cm, g, km, s
+from unyt.array import unyt_array
+
+
+def test_array_repr():
+    arr = [1, 2, 3] * cm
+    assert np.array_repr(arr) == "unyt_array([1, 2, 3], units='cm')"
+
+
+def test_dot_vectors():
+    a = [1, 2, 3] * cm
+    b = [1, 2, 3] * s
+    res = np.dot(a, b)
+    assert res.units == cm * s
+    assert res.d == 14
+
+
+# NOTE: explicitly setting the dtype of out arrays as `dtype=np.int_` is the
+# cross-platform way to guarantee that their dtype matches that of np.arange(x)
+# (on POSIX system it's int64, while on windows it's int32)
+@pytest.mark.parametrize(
+    "out",
+    [
+        None,
+        np.empty((3, 3), dtype=np.int_),
+        np.empty((3, 3), dtype=np.int_, order="C") * cm * s,
+        np.empty((3, 3), dtype=np.int_, order="C") * km * s,
+    ],
+    ids=[
+        "None",
+        "pure ndarray",
+        "same units",
+        "convertible units",
+    ],
+)
+def test_dot_matrices(out):
+    a = np.arange(9) * cm
+    a.shape = (3, 3)
+    b = np.arange(9) * s
+    b.shape = (3, 3)
+
+    res = np.dot(a, b, out=out)
+
+    if out is not None:
+        np.testing.assert_array_equal(res, out)
+        assert res is out
+
+    if isinstance(out, unyt_array):
+        # check that the result can be converted to predictible units
+        res.in_units("cm * s")
+        assert out.units == res.units
+
+
+def test_invalid_dot_matrices():
+    a = np.arange(9) * cm
+    a.shape = (3, 3)
+    b = np.arange(9) * s
+    b.shape = (3, 3)
+
+    out = np.empty((3, 3), dtype=np.int_, order="C") * s**2
+    with pytest.raises(
+        TypeError,
+        match=re.escape(
+            "output array is not acceptable "
+            "(units 's**2' cannot be converted to 'cm*s')"
+        ),
+    ):
+        np.dot(a, b, out=out)
+
+
+def test_vdot():
+    a = np.arange(9) * cm
+    b = np.arange(9) * s
+    res = np.vdot(a, b)
+    assert res.units == cm * s
+
+
+def test_inner():
+    a = np.array([1, 2, 3]) * cm
+    b = np.array([0, 1, 0]) * s
+    res = np.inner(a, b)
+    assert res.d == 2
+    assert res.units == cm * s
+
+
+def test_outer():
+    a = np.array([1, 2, 3]) * cm
+    b = np.array([0, 1, 0]) * s
+    res = np.outer(a, b)
+    expected = np.array([[0, 1, 0], [0, 2, 0], [0, 3, 0]])
+    np.testing.assert_array_equal(res.ndview, expected)
+    assert res.units == cm * s
+
+
+def test_kron():
+    a = np.eye(2) * cm
+    b = np.ones((2, 2)) * s
+    res = np.kron(a, b)
+    assert res.units == cm * s
+
+
+def test_matmul():
+    a = np.array([[1, 0], [0, 1]]) * cm
+    b = np.array([[4, 1], [2, 2]]) * s
+    expected = np.array([[4, 1], [2, 2]]) * cm * s
+    out = np.empty_like(expected)
+    res = np.matmul(a, b, out=out)
+    np.testing.assert_array_equal(res, expected)
+    assert res.units == expected.units
+
+    with pytest.xfail(
+        reason=(
+            "At of numpy 1.21.2, np.matmul seem to "
+            "escape the __unyt_array__ protocol"
+        )
+    ):
+        assert res is out
+
+
+def test_linalg_inv():
+    arr = np.random.random_sample((3, 3)) * cm
+    iarr = np.linalg.inv(arr)
+    assert 1 * iarr.units == 1 / cm
+
+
+def test_linalg_tensorinv():
+    a = np.eye(4 * 6) * cm
+    a.shape = (4, 6, 8, 3)
+    ia = np.linalg.tensorinv(a)
+    assert 1 * ia.units == 1 / cm
+
+
+def test_linalg_pinv():
+    a = np.random.randn(9, 6) * cm
+    B = np.linalg.pinv(a)
+    assert 1 * B.units == 1 / cm
+    np.testing.assert_allclose(a, np.dot(a, np.dot(B, a)))
+    np.testing.assert_allclose(B, np.dot(B, np.dot(a, B)))
+
+
+@pytest.mark.xfail(
+    reason=(
+        "as of numpy 1.21.2, the __array_function__ protocol doesn't let "
+        "us overload np.linalg.pinv when the first argument isn't a unyt_array"
+    )
+)
+def test_matrix_stack_linalg_pinv():
+    stack = [np.eye(4) * g for _ in range(3)]
+    B = np.linalg.pinv(stack)
+    assert 1 * B.units == 1 / g
+
+
+@pytest.mark.xfail(
+    reason=(
+        "as of numpy 1.21.2, the __array_function__ protocol doesn't let "
+        "us overload np.linalg.pinv when the first argument isn't a unyt_array"
+    )
+)
+def test_invalid_matrix_stack_linalg_pinv():
+    stack = [np.eye(4) * g, np.eye(4) * s]
+    with pytest.raises(
+        TypeError,
+        match=re.escape(
+            "numpy.linalg.pinv cannot operate on a stack "
+            "of matrices with different units."
+        ),
+    ):
+        np.linalg.pinv(stack)
+
+
+def test_histogram():
+    arr = np.random.normal(size=1000) * cm
+    counts, bins = np.histogram(arr, bins=10, range=(arr.min(), arr.max()))
+    assert type(counts) is np.ndarray
+    assert bins.units == arr.units
+
+
+def test_histogram2d():
+    x = np.random.normal(size=100) * cm
+    y = np.random.normal(loc=10, size=100) * s
+    counts, xbins, ybins = np.histogram2d(x, y)
+    assert counts.ndim == 2
+    assert 1 * xbins.units == 1 * x.units
+    assert 1 * ybins.units == 1 * y.units
+
+
+def test_histogramdd():
+    x = np.random.normal(size=100) * cm
+    y = np.random.normal(size=100) * s
+    z = np.random.normal(size=100) * g
+    counts, (xbins, ybins, zbins) = np.histogramdd((x, y, z))
+    assert counts.ndim == 3
+    assert 1 * xbins.units == 1 * x.units
+    assert 1 * ybins.units == 1 * y.units
+    assert 1 * zbins.units == 1 * z.units
diff --git a/unyt/tests/test_unyt_array.py b/unyt/tests/test_unyt_array.py
index 5df8020..6f8558f 100644
--- a/unyt/tests/test_unyt_array.py
+++ b/unyt/tests/test_unyt_array.py
@@ -2064,7 +2064,8 @@ def test_numpy_wrappers():
     assert_array_equal(unorm(a3, axis=1), unyt_array(arr_norm_answer, "cm"))
     assert_array_equal(np.linalg.norm(a3, axis=1), arr_norm_answer)
 
-    assert_array_equal(udot(a1, a1), unyt_quantity(dot_answer, "cm**2"))
+    with pytest.warns(DeprecationWarning):
+        assert_array_equal(udot(a1, a1), unyt_quantity(dot_answer, "cm**2"))
 
     assert_array_equal(np.array(catenate_answer), uconcatenate((a1.v, a2.v)))
     with pytest.raises(RuntimeError):