Repository: yt-project/unyt

Issue:
An error occurs when the exponent is an array
* unyt version: 2.9.5
* Python version: 3.9.17
* Operating System: 14.6.1 

### Description

In the `power` operation of unyt arrays, when the exponent is also an array, it leads to an error.

```shell
In [1]: from unyt import unyt_array

In [2]: a = unyt_array([1,2,3,4])

In [3]: b = unyt_array([1,2,3,4])

In [4]: a ** b
---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
Cell In[4], line 1
----> 1 a ** b

File ~/Apps/anaconda3/lib/python3.9/site-packages/unyt/array.py:1745, in unyt_array.__pow__(self, p, mod)
   1737 def __pow__(self, p, mod=None, /):
   1738     """
   1739     Power function, over-rides the ufunc as
   1740     ``numpy`` passes the ``_ones_like`` ufunc for
   1741     powers of zero (in some cases), which leads to an
   1742     incorrect unit calculation.
   1743     """
-> 1745     if p == 0.0:
   1746         ret = self.ua
   1747         ret.units = Unit("dimensionless")

ValueError: The truth value of an array with more than one element is ambiguous. Use a.any() or a.all()
```

While for numpy arrays, it behaviors like 

```shell
In [5]: import numpy as np

In [6]: a = np.array([1,2,3,4])

In [7]: b = np.array([1,2,3,4])

In [8]: a ** b
Out[8]: array([  1,   4,  27, 256])
```

I now use `to_ndarray()` to make it work but lose the units.

Hints from issue comments:
@neutrinoceros: Thanks for reporting this, it's clearly a defect !
It might sufficient to replace `p == 0.0` with `np.all(p == 0.0)` and add a couple test cases. We should also make sure that a TypeError is raised if p is another (non dimensionless) unyt array. Are you interested in giving it a try ? Otherwise I should have time to do it myself later this week.@xshaokun: Thank you for the response. From the perspective of dimensions, `p` as an exponent here should not be a quantity with dimensions, so I don't think we need to worry about `p` being non-dimensionless. Since I will be quite busy for a while and may not be able to fix this bug, I appreciate you would make it :)@neutrinoceros: Hum, maybe I was a bit quick in my previous answer: the only cases where it makes sense to support raising a `unyt_array` `a` to an array power `p` is when `p` is a 1-element array or has a broadcastable shape with `a` and a unique value. The reason is that the resulting array can only have a uniform unit, which is only possible if all elements are raised to the same power... so, even though the specific example you illustrated the issue with doesn't make sense to me (but the error should be clearer), I'll add support for the few cases that *do* make sense.@xshaokun: Oh, in some cases, when the base is a dimensionless number, the example I gave should make sense. For instance, I’m currently trying to calculate synchrotron radiation, where the spectral indices vary across different regions. The following formula shows that the base in the final part is a dimensionless quantity, where the `p` is the power-law index of non-thermal electrons.

<img width="654" alt="image" src="https://github.com/user-attachments/assets/826ccfd2-aa17-442b-8450-35b66bd33e81">

(from Rybicki and Lightman 1986)@neutrinoceros: Yes, if the power itself is dimensionless, it should work. See #524 and please let me know if this helps !

Solution patch:
diff --git a/.github/workflows/bleeding-edge.yaml b/.github/workflows/bleeding-edge.yaml
index fd6cf91f..8c673ee7 100644
--- a/.github/workflows/bleeding-edge.yaml
+++ b/.github/workflows/bleeding-edge.yaml
@@ -12,9 +12,9 @@ on:
   pull_request:
     paths:
       - .github/workflows/bleeding-edge.yaml
-  # Run daily at 0:01 UTC
+  # Run every Wednesday at 0:01 UTC
   schedule:
-    - cron:  '1 0 * * *'
+    - cron:  1 0 * * 3
   workflow_dispatch:
 
 jobs:
diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml
index 6a833531..30a4e2fd 100644
--- a/.github/workflows/ci.yml
+++ b/.github/workflows/ci.yml
@@ -24,8 +24,8 @@ jobs:
           - '3.12'
         # Test all on ubuntu, test ends on macos and windows
         include:
-          - os: macos-12
-            # pin macos-12 (x86) because Python 3.9 is broken in the arm64 image
+          - os: macos-13
+            # pin macos-13 (x86) because Python 3.9 is broken in the arm64 image
             python-version: '3.9'
           - os: windows-latest
             python-version: '3.9'
diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
index 3a9a68f6..243e0ae1 100644
--- a/.pre-commit-config.yaml
+++ b/.pre-commit-config.yaml
@@ -20,20 +20,18 @@ repos:
           - id: check-yaml
           - id: check-toml
 
-    - repo: https://github.com/psf/black
-      rev: 24.10.0
-      hooks:
-          - id: black
-
+    # TODO: replace this with ruff when it supports embedded python blocks
+    # see https://github.com/astral-sh/ruff/issues/8237
     - repo: https://github.com/adamchainz/blacken-docs
       rev: 1.19.1
       hooks:
           - id: blacken-docs
-            additional_dependencies: [black==24.1.1]
+            additional_dependencies: [black==24.10.0]
 
     - repo: https://github.com/astral-sh/ruff-pre-commit
-      rev: v0.7.1
+      rev: v0.9.9
       hooks:
+          - id: ruff-format
           - id: ruff
             args: [--fix, --show-fixes]
 
@@ -43,7 +41,7 @@ repos:
           - id: rst-backticks
 
     - repo: https://github.com/codespell-project/codespell
-      rev: v2.3.0
+      rev: v2.4.1
       hooks:
           - id: codespell
             args: [--skip="*.pdf", -L, 'smoot,copin,celcius,ue,bu,dne,ond']
diff --git a/unyt/_array_functions.py b/unyt/_array_functions.py
index 646629e2..56283a84 100644
--- a/unyt/_array_functions.py
+++ b/unyt/_array_functions.py
@@ -158,38 +158,169 @@ def _sanitize_range(_range, units):
     return new_range.squeeze()
 
 
-@implements(np.histogram)
-def histogram(
-    a,
-    bins=10,
-    range=None,
-    *args,
-    **kwargs,
-):
-    range = _sanitize_range(range, units=[a.units])
-    counts, bins = np.histogram._implementation(
-        np.asarray(a), bins, range, *args, **kwargs
-    )
-    return counts, bins * a.units
+def _histogram(a, *, bins=10, range=None, density=None, weights=None, normed=None):
+    range = _sanitize_range(range, units=[getattr(a, "units", None)])
+    if NUMPY_VERSION >= Version("1.24"):
+        counts, bins = np.histogram._implementation(
+            np.asarray(a),
+            bins=bins,
+            range=range,
+            density=density,
+            weights=np.asarray(weights) if weights is not None else None,
+        )
+    else:
+        counts, bins = np.histogram._implementation(
+            np.asarray(a),
+            bins=bins,
+            range=range,
+            normed=normed,
+            weights=np.asarray(weights) if weights is not None else None,
+            density=density,
+        )
+    # a and/or weights could have units, only apply if present
+    # don't getattr(..., "units", NULL_UNIT) because e.g. we don't want
+    # a unyt_array if weights are not a unyt_array and not density
+    if density and hasattr(a, "units"):
+        counts /= a.units
+    if weights is not None and hasattr(weights, "units"):
+        counts *= weights.units
+    return counts, bins * getattr(a, "units", 1)
 
 
-@implements(np.histogram2d)
-def histogram2d(x, y, bins=10, range=None, *args, **kwargs):
-    range = _sanitize_range(range, units=[x.units, y.units])
-    counts, xbins, ybins = np.histogram2d._implementation(
-        np.asarray(x), np.asarray(y), bins, range, *args, **kwargs
-    )
-    return counts, xbins * x.units, ybins * y.units
+if NUMPY_VERSION >= Version("1.24"):
 
+    @implements(np.histogram)
+    def histogram(a, bins=10, range=None, density=None, weights=None):
+        return _histogram(a, bins=bins, range=range, density=density, weights=weights)
 
-@implements(np.histogramdd)
-def histogramdd(sample, bins=10, range=None, *args, **kwargs):
-    units = [_.units for _ in sample]
-    range = _sanitize_range(range, units=units)
-    counts, bins = np.histogramdd._implementation(
-        [np.asarray(_) for _ in sample], bins, range, *args, **kwargs
+else:
+
+    @implements(np.histogram)
+    def histogram(a, bins=10, range=None, normed=None, weights=None, density=None):
+        return _histogram(
+            a, bins=bins, range=range, normed=normed, weights=weights, density=density
+        )
+
+
+def _histogram2d(x, y, *, bins=10, range=None, density=None, weights=None, normed=None):
+    range = _sanitize_range(
+        range, units=[getattr(x, "units", None), getattr(y, "units", None)]
     )
-    return counts, tuple(_bin * u for _bin, u in zip(bins, units))
+    if NUMPY_VERSION >= Version("1.24"):
+        counts, xbins, ybins = np.histogram2d._implementation(
+            np.asarray(x),
+            np.asarray(y),
+            bins=bins,
+            range=range,
+            density=density,
+            weights=np.asarray(weights) if weights is not None else None,
+        )
+    else:
+        counts, xbins, ybins = np.histogram2d._implementation(
+            np.asarray(x),
+            np.asarray(y),
+            bins=bins,
+            range=range,
+            normed=normed,
+            weights=np.asarray(weights) if weights is not None else None,
+            density=density,
+        )
+    # x, y and/or weights could have units, only apply if present
+    # don't getattr(..., "units", NULL_UNIT) because e.g. we don't want
+    # a unyt_array if weights are not a unyt_array and not density
+    if density:
+        if hasattr(x, "units"):
+            counts /= x.units
+        if hasattr(y, "units"):
+            counts /= y.units
+    if weights is not None and hasattr(weights, "units"):
+        counts *= weights.units
+    return counts, xbins * getattr(x, "units", 1), ybins * getattr(y, "units", 1)
+
+
+if NUMPY_VERSION >= Version("1.24"):
+
+    @implements(np.histogram2d)
+    def histogram2d(x, y, bins=10, range=None, density=None, weights=None):
+        return _histogram2d(
+            x, y, bins=bins, range=range, density=density, weights=weights
+        )
+
+else:
+
+    @implements(np.histogram2d)
+    def histogram2d(x, y, bins=10, range=None, normed=None, weights=None, density=None):
+        return _histogram2d(
+            x,
+            y,
+            bins=bins,
+            range=range,
+            normed=normed,
+            weights=weights,
+            density=density,
+        )
+
+
+def _histogramdd(
+    sample, *, bins=10, range=None, density=None, weights=None, normed=None
+):
+    range = _sanitize_range(range, units=[getattr(_, "units", None) for _ in sample])
+    if NUMPY_VERSION >= Version("1.24"):
+        counts, bins = np.histogramdd._implementation(
+            [np.asarray(_) for _ in sample],
+            bins=bins,
+            range=range,
+            density=density,
+            weights=np.asarray(weights) if weights is not None else None,
+        )
+    else:
+        counts, bins = np.histogramdd._implementation(
+            [np.asarray(_) for _ in sample],
+            bins=bins,
+            range=range,
+            normed=normed,
+            weights=np.asarray(weights) if weights is not None else None,
+            density=density,
+        )
+    # sample(s) and/or weights could have units, only apply if present
+    # don't getattr(..., "units", NULL_UNIT) because e.g. we don't want
+    # a unyt_array if weights are not a unyt_array and not density
+    if density:
+        divider_units = 1 * NULL_UNIT
+        for s in sample:
+            if not hasattr(s, "units"):
+                continue
+            divider_units *= s.units
+        if divider_units != NULL_UNIT:
+            counts /= divider_units
+
+    if weights is not None and hasattr(weights, "units"):
+        counts *= weights.units
+    return counts, tuple(_bin * getattr(s, "units", 1) for _bin, s in zip(bins, sample))
+
+
+if NUMPY_VERSION >= Version("1.24"):
+
+    @implements(np.histogramdd)
+    def histogramdd(sample, bins=10, range=None, density=None, weights=None):
+        return _histogramdd(
+            sample, bins=bins, range=range, density=density, weights=weights
+        )
+
+else:
+
+    @implements(np.histogramdd)
+    def histogramdd(
+        sample, bins=10, range=None, normed=None, weights=None, density=None
+    ):
+        return _histogramdd(
+            sample,
+            bins=bins,
+            range=range,
+            normed=normed,
+            weights=weights,
+            density=density,
+        )
 
 
 @implements(np.histogram_bin_edges)
@@ -490,26 +621,94 @@ def array_equiv(a1, a2, *args, **kwargs) -> bool:
     )
 
 
-@implements(np.linspace)
-def linspace(start, stop, *args, **kwargs):
+def _linspace(
+    start,
+    stop,
+    num=50,
+    endpoint=True,
+    retstep=False,
+    dtype=None,
+    axis=0,
+    *,
+    device=None,
+):
     _validate_units_consistency((start, stop))
-    return (
-        np.linspace._implementation(
-            np.asarray(start), np.asarray(stop), *args, **kwargs
+    kwargs = {
+        "num": num,
+        "endpoint": endpoint,
+        "retstep": retstep,
+        "dtype": dtype,
+        "axis": axis,
+    }
+    if NUMPY_VERSION >= Version("2.1.0.dev0"):
+        kwargs["device"] = device
+    result = np.linspace._implementation(np.asarray(start), np.asarray(stop), **kwargs)
+    if retstep:
+        return result[0] * start.units, result[1] * start.units
+    else:
+        return result * start.units
+
+
+if NUMPY_VERSION >= Version("2.0.0.dev0"):
+
+    @implements(np.linspace)
+    def linspace(
+        start,
+        stop,
+        num=50,
+        endpoint=True,
+        retstep=False,
+        dtype=None,
+        axis=0,
+        *,
+        device=None,
+    ):
+        return _linspace(
+            start,
+            stop,
+            num=num,
+            endpoint=endpoint,
+            retstep=retstep,
+            dtype=dtype,
+            axis=axis,
+            device=device,
+        )
+
+else:
+
+    @implements(np.linspace)
+    def linspace(start, stop, num=50, endpoint=True, retstep=False, dtype=None, axis=0):
+        return _linspace(
+            start,
+            stop,
+            num=num,
+            endpoint=endpoint,
+            retstep=retstep,
+            dtype=dtype,
+            axis=axis,
         )
-        * start.units
-    )
 
 
 @implements(np.logspace)
-def logspace(start, stop, *args, **kwargs):
-    _validate_units_consistency((start, stop))
-    return (
-        np.logspace._implementation(
-            np.asarray(start), np.asarray(stop), *args, **kwargs
+def logspace(start, stop, num=50, endpoint=True, base=10, dtype=None, axis=0):
+    startu = getattr(start, "units", NULL_UNIT)
+    stopu = getattr(stop, "units", NULL_UNIT)
+    if startu != NULL_UNIT or stopu != NULL_UNIT:
+        raise TypeError(
+            "The first two arguments to numpy.logspace must be dimensionless, "
+            f"got units={startu} (arg1) and units={stopu} (arg2). If output with"
+            " units is desired, apply the units to the `base` kwarg."
         )
-        * start.units
+    result = np.logspace._implementation(
+        np.asarray(start),
+        np.asarray(stop),
+        num=num,
+        endpoint=endpoint,
+        base=np.asarray(base),
+        dtype=dtype,
+        axis=axis,
     )
+    return result * getattr(base, "units", NULL_UNIT)
 
 
 @implements(np.geomspace)
@@ -535,7 +734,8 @@ def copyto(dst, src, *args, **kwargs):
 
 @implements(np.prod)
 def prod(a, *args, **kwargs):
-    return np.prod._implementation(np.asarray(a), *args, **kwargs) * a.units**a.size
+    res = np.prod._implementation(np.asarray(a), *args, **kwargs)
+    return res * a.units ** (a.size // res.size)
 
 
 @implements(np.var)
@@ -570,8 +770,9 @@ def nanquantile(a, *args, **kwargs):
 
 @implements(np.linalg.det)
 def linalg_det(a, *args, **kwargs):
-    return np.linalg.det._implementation(np.asarray(a), *args, **kwargs) * a.units ** (
-        a.shape[0]
+    return (
+        np.linalg.det._implementation(np.asarray(a), *args, **kwargs)
+        * a.units ** (a.shape[0])
     )
 
 
@@ -649,7 +850,7 @@ def savetxt(fname, X, *args, **kwargs):
 
 @implements(np.apply_over_axes)
 def apply_over_axes(func, a, axes):
-    res = func(np.asarray(a), axes[0]) * a.units
+    res = func(a, axes[0])
     if len(axes) > 1:
         # this function is recursive by nature,
         # here we intentionally do not call the base _implementation
@@ -1019,3 +1220,23 @@ def in1d(ar1, ar2, *args, **kwargs):
         return np.isin._implementation(
             np.asarray(ar1), np.asarray(ar2), *args, **kwargs
         )
+
+
+@implements(np.take)
+def take(a, indices, axis=None, out=None, mode="raise"):
+    ret_units = getattr(a, "units", NULL_UNIT)
+
+    if out is not None:
+        out_view = np.asarray(out)
+    else:
+        out_view = None
+
+    res = np.take._implementation(
+        np.asarray(a), indices, axis=axis, out=out_view, mode=mode
+    )
+
+    if getattr(out, "units", None) is not None:
+        out.units = ret_units
+
+    ret_cls = unyt_quantity if res.ndim == 0 else unyt_array
+    return ret_cls(res, ret_units, bypass_validation=True)
diff --git a/unyt/array.py b/unyt/array.py
index 330f51c4..b1944119 100644
--- a/unyt/array.py
+++ b/unyt/array.py
@@ -132,6 +132,10 @@
 
 from ._deprecation import warn_deprecated
 
+NUMPY_VERSION = Version(version("numpy"))
+if NUMPY_VERSION >= Version("2.0.0.dev0"):
+    from numpy import vecdot
+
 NULL_UNIT = Unit()
 POWER_MAPPING = {multiply: lambda x: x, divide: lambda x: 2 - x}
 DISALLOWED_DTYPES = (
@@ -424,6 +428,8 @@ def _apply_power_mapping(ufunc, in_unit, in_size, in_shape, input_kwarg_dict):
     divmod_,
     heaviside,
 )
+if NUMPY_VERSION >= Version("2.0.0.dev0"):
+    binary_operators += (vecdot,)
 
 trigonometric_operators = (sin, cos, tan)
 
@@ -576,6 +582,8 @@ class unyt_array(np.ndarray):
         matmul: _multiply_units,
         clip: _passthrough_unit,
     }
+    if NUMPY_VERSION >= Version("2.0.0.dev0"):
+        _ufunc_registry[vecdot] = _multiply_units
 
     __array_priority__ = 2.0
 
@@ -1764,8 +1772,8 @@ def __pow__(self, p, mod=None, /):
         Power function
         """
         # see https://github.com/yt-project/unyt/issues/203
-        if p == 0.0:
-            ret = self.ua
+        if np.isscalar(p) and p == 0.0:
+            ret = self.unit_array
             ret.units = Unit("dimensionless")
             return ret
         else:
@@ -1854,17 +1862,32 @@ def __array_ufunc__(self, ufunc, method, *inputs, **kwargs):
                 u1 = Unit(registry=getattr(u0, "registry", None))
             elif ufunc is power:
                 u1 = inp1
-                if inp0.shape != () and inp1.shape != ():
-                    raise UnitOperationError(ufunc, u0, u1)
-                if isinstance(u1, unyt_array):
-                    if u1.units.is_dimensionless:
-                        pass
-                    else:
+                if inp0.shape == () or inp1.shape == ():
+                    if isinstance(u1, unyt_array) and not u1.units.is_dimensionless:
                         raise UnitOperationError(ufunc, u0, u1.units)
-                if u1.shape == ():
-                    u1 = float(u1)
+                    if u1.shape == ():
+                        u1 = float(u1)
+                    else:
+                        u1 = 1.0
+                elif inp0.shape == inp1.shape:
+                    if isinstance(u1, unyt_array) and not u1.units.is_dimensionless:
+                        raise UnitOperationError(ufunc, u0, getattr(u1, "units", None))
+
+                    if (
+                        (isinstance(u0, Unit) and not u0.is_dimensionless)
+                        or isinstance(u0, unyt_array)
+                        and not u0.units.is_dimensionless
+                    ):
+                        # u0 has units
+                        if np.ptp(u1) != 0:
+                            raise UnitOperationError(
+                                ufunc, u0, getattr(u1, "units", None)
+                            )
+
+                    first_element_slice = (0,) * u1.ndim
+                    u1 = float(u1[first_element_slice])
                 else:
-                    u1 = 1.0
+                    raise UnitOperationError(ufunc, u0, u1)
             unit_operator = self._ufunc_registry[ufunc]
 
             if (
@@ -1985,8 +2008,8 @@ def __array_ufunc__(self, ufunc, method, *inputs, **kwargs):
                 mul = 1
             else:
                 raise RuntimeError(
-                    "Support for the %s ufunc with %i inputs has not been "
-                    "added to unyt_array." % (str(ufunc), len(inputs))
+                    f"Support for the {ufunc} ufunc with {len(inputs)} "
+                    "inputs has not been added to unyt_array."
                 )
         if unit is None:
             out_arr = np.array(out_arr, copy=_COPY_IF_NEEDED)
@@ -1997,7 +2020,7 @@ def __array_ufunc__(self, ufunc, method, *inputs, **kwargs):
         elif out_arr.size == 1:
             out_arr = unyt_array(np.asarray(out_arr), unit)
         else:
-            if ret_class is unyt_quantity:
+            if issubclass(ret_class, unyt_quantity):
                 # This happens if you do ndarray * unyt_quantity.
                 # Explicitly casting to unyt_array avoids creating a
                 # unyt_quantity with size > 1
@@ -2125,6 +2148,21 @@ def dot(self, b, out=None):
             out.units = res_units
         return ret
 
+    def take(self, indices, axis=None, out=None, mode="raise"):
+        """method
+
+        Return an array formed from the elements of `a` at the given indices.
+
+        Refer to :func:`numpy.take` for full documentation.
+
+        See also
+        --------
+        numpy.take : equivalent function
+        """
+        from ._array_functions import take
+
+        return take(self, indices, axis=axis, out=out, mode=mode)
+
     def __reduce__(self):
         """Pickle reduction method
 
diff --git a/unyt/exceptions.py b/unyt/exceptions.py
index 93ecb860..750511d2 100644
--- a/unyt/exceptions.py
+++ b/unyt/exceptions.py
@@ -251,8 +251,7 @@ def __str__(self):
 
         if isinstance(self.unit2, Unit):
             msg = (
-                "The unit equivalence '%s' does not exist for the units "
-                "'%s' and '%s'."
+                "The unit equivalence '%s' does not exist for the units '%s' and '%s'."
             )
         else:
             msg = (
diff --git a/unyt/unit_object.py b/unyt/unit_object.py
index bec72b2c..0ffe7284 100644
--- a/unyt/unit_object.py
+++ b/unyt/unit_object.py
@@ -442,8 +442,7 @@ def __truediv__(self, u):
                 base_offset = self.base_offset
             else:
                 raise InvalidUnitOperation(
-                    "Quantities with units of Farhenheit "
-                    "and Celsius cannot be divided."
+                    "Quantities with units of Farhenheit and Celsius cannot be divided."
                 )
 
         return Unit(
@@ -1070,9 +1069,7 @@ def define_unit(
         if _iterable(value) and len(value) == 2:
             value = unyt_quantity(value[0], value[1], registry=registry)
         else:
-            raise RuntimeError(
-                '"value" needs to be a quantity or ' "(value, unit) tuple!"
-            )
+            raise RuntimeError('"value" needs to be a quantity or (value, unit) tuple!')
     base_value = float(value.in_base(unit_system="mks"))
     dimensions = value.units.dimensions
     registry.add(


Test patch:
diff --git a/unyt/tests/test_array_functions.py b/unyt/tests/test_array_functions.py
index bfd900ab..bcff8e6a 100644
--- a/unyt/tests/test_array_functions.py
+++ b/unyt/tests/test_array_functions.py
@@ -9,7 +9,7 @@
 from numpy.testing import assert_allclose
 from packaging.version import Version
 
-from unyt import A, K, cm, degC, delta_degC, g, km, rad, s
+from unyt import A, K, Msun, cm, degC, delta_degC, dimensionless, g, km, rad, s
 from unyt._array_functions import (
     _HANDLED_FUNCTIONS as HANDLED_FUNCTIONS,
     _UNSUPPORTED_FUNCTIONS as UNSUPPORTED_FUNCTIONS,
@@ -143,7 +143,6 @@
     np.nancumprod,  # we get it for free with np.cumprod (tested)
     np.bincount,  # works out of the box (tested)
     np.unique,  # works out of the box (tested)
-    np.take,  # works out of the box (tested)
     np.min_scalar_type,  # returns dtypes
     np.extract,  # works out of the box (tested)
     np.setxor1d,  # we get it for free with previously implemented functions (tested)
@@ -181,7 +180,6 @@
         np.unique_counts,
         np.unique_inverse,
         np.unique_values,
-        np.vecdot,
     }
 
 if NUMPY_VERSION >= Version("2.1.0dev0"):
@@ -296,9 +294,9 @@ def test_wrapping_completeness():
     """Ensure we wrap all numpy functions that support __array_function__"""
     handled_numpy_functions = set(HANDLED_FUNCTIONS.keys())
     # ensure no functions appear in both NOT_HANDLED_FUNCTIONS and HANDLED_FUNCTIONS
-    assert NOT_HANDLED_FUNCTIONS.isdisjoint(
-        handled_numpy_functions
-    ), NOT_HANDLED_FUNCTIONS.intersection(handled_numpy_functions)
+    assert NOT_HANDLED_FUNCTIONS.isdisjoint(handled_numpy_functions), (
+        NOT_HANDLED_FUNCTIONS.intersection(handled_numpy_functions)
+    )
     # get list of functions that support wrapping by introspection on numpy module
     wrappable_functions = get_wrapped_functions(np, np.fft, np.linalg)
     for function in HANDLED_FUNCTIONS:
@@ -589,16 +587,11 @@ def test_matrix_transpose(namespace):
 
 
 @pytest.mark.skipif(
-    NUMPY_VERSION < Version("2.0.0dev0"), reason="vecdot is new in numpy 2.0"
+    NUMPY_VERSION < Version("2.0.0dev0"), reason="linalg.vecdot is new in numpy 2.0"
 )
-@pytest.mark.parametrize("namespace", [None, "linalg"])
-def test_vecdot(namespace):
-    if namespace is None:
-        func = np.vecdot
-    else:
-        func = getattr(np, namespace).vecdot
+def test_linalg_vecdot():
     a = np.arange(0, 9)
-    assert_array_equal_units(func(a, a), np.vdot(a, a))
+    assert_array_equal_units(np.linalg.vecdot(a, a), np.vdot(a, a))
 
 
 @pytest.mark.skipif(
@@ -646,51 +639,203 @@ def test_linalg_tensordot():
     assert_array_equal_units(res, ref)
 
 
-def test_histogram():
-    rng = np.random.default_rng()
-    arr = rng.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_histogram_implicit_units():
-    # see https://github.com/yt-project/unyt/issues/465
-    rng = np.random.default_rng()
-    arr = rng.normal(size=1000) * cm
-    counts, bins = np.histogram(arr, bins=10, range=(arr.min().value, arr.max().value))
-    assert type(counts) is np.ndarray
-    assert bins.units == arr.units
-
-
-def test_histogram2d():
-    rng = np.random.default_rng()
-    x = rng.normal(size=100) * cm
-    y = rng.normal(loc=10, size=100) * s
-    counts, xbins, ybins = np.histogram2d(x, y)
-    assert counts.ndim == 2
-    assert xbins.units == x.units
-    assert ybins.units == y.units
-
-
-def test_histogramdd():
-    rng = np.random.default_rng()
-    x = rng.normal(size=100) * cm
-    y = rng.normal(size=100) * s
-    z = rng.normal(size=100) * g
-    counts, (xbins, ybins, zbins) = np.histogramdd((x, y, z))
-    assert counts.ndim == 3
-    assert xbins.units == x.units
-    assert ybins.units == y.units
-    assert zbins.units == z.units
-
-
-def test_histogram_bin_edges():
-    rng = np.random.default_rng()
-    arr = rng.normal(size=1000) * cm
-    bins = np.histogram_bin_edges(arr)
-    assert type(bins) is unyt_array
-    assert bins.units == arr.units
+class TestHistograms:
+    def test_histogram(self):
+        rng = np.random.default_rng()
+        arr = rng.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_histogram_implicit_units(self):
+        # see https://github.com/yt-project/unyt/issues/465
+        rng = np.random.default_rng()
+        arr = rng.normal(size=1000) * cm
+        counts, bins = np.histogram(
+            arr, bins=10, range=(arr.min().value, arr.max().value)
+        )
+        assert type(counts) is np.ndarray
+        assert bins.units == arr.units
+
+    def test_histogram_with_density(self):
+        rng = np.random.default_rng()
+        arr = rng.normal(size=1000) * cm
+        density, bins = np.histogram(
+            arr, bins=10, range=(arr.min(), arr.max()), density=True
+        )
+        assert type(density) is unyt_array
+        assert density.units == arr.units**-1
+        assert bins.units == arr.units
+
+    def test_histogram_with_weights(self):
+        rng = np.random.default_rng()
+        arr = rng.normal(size=1000) * cm
+        w = rng.uniform(size=1000) * g
+        wcounts, wbins = np.histogram(
+            arr, bins=10, range=(arr.min(), arr.max()), weights=w
+        )
+        assert type(wcounts) is unyt_array
+        assert wcounts.units == w.units
+        assert wbins.units == arr.units
+
+    def test_histogram_with_weights_and_density(self):
+        rng = np.random.default_rng()
+        arr = rng.normal(size=1000) * cm
+        w = rng.uniform(size=1000) * g
+        wdensity, wdbins = np.histogram(
+            arr, bins=10, range=(arr.min(), arr.max()), density=True, weights=w
+        )
+        assert type(wdensity) is unyt_array
+        assert wdensity.units == w.units / arr.units
+        assert wdbins.units == arr.units
+
+    def test_histogram_with_weights_and_dimless_arr(self):
+        rng = np.random.default_rng()
+        arr = rng.normal(size=1000) * cm
+        w = rng.uniform(size=1000) * g
+        wcounts2, wbins2 = np.histogram(arr.to_value(arr.units), weights=w)
+        assert type(wcounts2) is unyt_array
+        assert wcounts2.units == w.units
+        assert not hasattr(wbins2, "units")
+
+    def test_histogram2d(self):
+        rng = np.random.default_rng()
+        x = rng.normal(size=100) * cm
+        y = rng.normal(loc=10, size=100) * s
+        counts, xbins, ybins = np.histogram2d(x, y)
+        assert counts.ndim == 2
+        assert xbins.units == x.units
+        assert ybins.units == y.units
+
+    def test_histogram2d_with_density(self):
+        rng = np.random.default_rng()
+        x = rng.normal(size=100) * cm
+        y = rng.normal(loc=10, size=100) * s
+        density, xbins, ybins = np.histogram2d(x, y, density=True)
+        assert density.ndim == 2
+        assert type(density) is unyt_array
+        assert density.units == (x.units * y.units) ** -1
+        assert xbins.units == x.units
+        assert ybins.units == y.units
+
+    def test_histogram2d_with_weights(self):
+        rng = np.random.default_rng()
+        x = rng.normal(size=100) * cm
+        y = rng.normal(loc=10, size=100) * s
+        w = rng.uniform(size=100) * g
+        wcounts, xwbins, ywbins = np.histogram2d(x, y, weights=w)
+        assert wcounts.ndim == 2
+        assert type(wcounts) is unyt_array
+        assert wcounts.units == w.units
+        assert xwbins.units == x.units
+        assert ywbins.units == y.units
+
+    def test_histogram2d_with_weights_and_density(self):
+        rng = np.random.default_rng()
+        x = rng.normal(size=100) * cm
+        y = rng.normal(loc=10, size=100) * s
+        w = rng.uniform(size=100) * g
+        wdensity, xwdbins, ywdbins = np.histogram2d(x, y, weights=w, density=True)
+        assert wdensity.ndim == 2
+        assert type(wdensity) is unyt_array
+        assert wdensity.units == w.units / (x.units * y.units)
+        assert xwdbins.units == x.units
+        assert ywdbins.units == y.units
+
+    def test_histogram2d_with_weights_and_dimless_arr(self):
+        rng = np.random.default_rng()
+        x = rng.normal(size=100) * cm
+        y = rng.normal(loc=10, size=100) * s
+        w = rng.uniform(size=100) * g
+        wcounts2, xwbins2, ywbins2 = np.histogram2d(
+            x.to_value(x.units), y.to_value(y.units), weights=w
+        )
+        assert type(wcounts2) is unyt_array
+        assert wcounts2.units == w.units
+        assert not hasattr(xwbins2, "units")
+        assert not hasattr(ywbins2, "units")
+
+    def test_histogramdd(self):
+        rng = np.random.default_rng()
+        x = rng.normal(size=100) * cm
+        y = rng.normal(size=100) * s
+        z = rng.normal(size=100) * g
+        counts, (xbins, ybins, zbins) = np.histogramdd((x, y, z))
+        assert counts.ndim == 3
+        assert xbins.units == x.units
+        assert ybins.units == y.units
+        assert zbins.units == z.units
+
+    def test_histogramdd_with_density(self):
+        rng = np.random.default_rng()
+        x = rng.normal(size=100) * cm
+        y = rng.normal(loc=10, size=100) * s
+        z = rng.normal(size=100) * g
+        density, (xbins, ybins, zbins) = np.histogramdd((x, y, z), density=True)
+        assert density.ndim == 3
+        assert type(density) is unyt_array
+        assert density.units == (x.units * y.units * z.units) ** -1
+        assert xbins.units == x.units
+        assert ybins.units == y.units
+        assert zbins.units == z.units
+
+    def test_histogramdd_with_weights(self):
+        rng = np.random.default_rng()
+        x = rng.normal(size=100) * cm
+        y = rng.normal(loc=10, size=100) * s
+        z = rng.normal(size=100) * g
+        w = rng.uniform(size=100) * K
+        wcounts, (xwbins, ywbins, zwbins) = np.histogramdd((x, y, z), weights=w)
+        assert wcounts.ndim == 3
+        assert type(wcounts) is unyt_array
+        assert wcounts.units == w.units
+        assert xwbins.units == x.units
+        assert ywbins.units == y.units
+        assert zwbins.units == z.units
+
+    def test_histogramdd_with_weights_and_density(self):
+        rng = np.random.default_rng()
+        x = rng.normal(size=100) * cm
+        y = rng.normal(loc=10, size=100) * s
+        z = rng.normal(size=100) * g
+        w = rng.uniform(size=100) * K
+        wdensity, (xwdbins, ywdbins, zwdbins) = np.histogramdd(
+            (x, y, z), weights=w, density=True
+        )
+        assert wdensity.ndim == 3
+        assert type(wdensity) is unyt_array
+        assert wdensity.units == w.units / (x.units * y.units * z.units)
+        assert xwdbins.units == x.units
+        assert ywdbins.units == y.units
+        assert zwdbins.units == z.units
+
+    def test_histogramdd_with_weights_and_dimless_arr(self):
+        rng = np.random.default_rng()
+        x = rng.normal(size=100) * cm
+        y = rng.normal(loc=10, size=100) * s
+        z = rng.normal(size=100) * g
+        w = rng.uniform(size=100) * K
+        wcounts2, (xwbins2, ywbins2, zwbins2) = np.histogramdd(
+            (x.to_value(x.units), y.to_value(y.units), z.to_value(z.units)), weights=w
+        )
+        assert type(wcounts2) is unyt_array
+        assert wcounts2.units == w.units
+        assert not hasattr(xwbins2, "units")
+        assert not hasattr(ywbins2, "units")
+        assert not hasattr(zwbins2, "units")
+
+    @pytest.mark.parametrize("weights", [None, [0, 1, 2], [0, 1, 2] * cm])
+    def test_histogramdd_recursion(self, weights):
+        # regression test for https://github.com/yt-project/unyt/issues/540
+        sample = [unyt_array(np.arange(3), Msun)]
+        np.histogramdd(sample, density=True, weights=weights)
+
+    def test_histogram_bin_edges(self):
+        rng = np.random.default_rng()
+        arr = rng.normal(size=1000) * cm
+        bins = np.histogram_bin_edges(arr)
+        assert type(bins) is unyt_array
+        assert bins.units == arr.units
 
 
 def test_concatenate():
@@ -1081,16 +1226,42 @@ def test_isclose_error():
         np.isclose(x, y)
 
 
-@pytest.mark.parametrize(
-    "func",
-    [
-        np.linspace,
-        np.logspace,
-        np.geomspace,
-    ],
-)
-def test_xspace(func):
-    res = func(1 * cm, 11 * cm, 10)
+def test_linspace():
+    res = np.linspace(1 * cm, 11 * cm, 10)
+    assert type(res) is unyt_array
+    assert res.units == cm
+
+
+def test_linspace_with_retstep():
+    res, step = np.linspace(1 * cm, 11 * cm, 10, retstep=True)
+    assert type(res) is unyt_array
+    assert res.units == cm
+    assert type(step) is unyt_quantity
+    assert step.units == cm
+
+
+def test_logspace_with_units_raises():
+    with pytest.raises(
+        TypeError,
+        match="The first two arguments to numpy.logspace must be dimensionless",
+    ):
+        np.logspace(1 * cm, 2 * cm)
+
+
+def test_logspace_with_base():
+    res = np.logspace(1, 2, base=10 * cm)
+    assert type(res) is unyt_array
+    assert res.units == cm
+
+
+def test_logspace_with_dimless_input():
+    res = np.logspace(1 * dimensionless, 2 * dimensionless)
+    assert type(res) is unyt_array
+    assert res.units == dimensionless
+
+
+def test_geomspace():
+    res = np.geomspace(1 * cm, 11 * cm, 10)
     assert type(res) is unyt_array
     assert res.units == cm
 
@@ -1291,6 +1462,18 @@ def test_scalar_reducer(func, expected_units):
     assert y.units == expected_units
 
 
+@pytest.mark.parametrize("func", [np.prod, np.nanprod])
+def test_prod_with_axis(func):
+    x = [
+        [1, 2, 3],
+        [4, 5, 6],
+        [7, 8, 9],
+    ] * cm
+    y = func(x, axis=0)
+    assert type(y) is unyt_array
+    assert y.units == cm**3
+
+
 @pytest.mark.parametrize(
     "func",
     [
@@ -1464,11 +1647,16 @@ def test_apply_along_axis():
     assert ret.units == cm**2
 
 
-def test_apply_over_axes():
+@pytest.mark.parametrize("axes, expected_units", [((0, 1), cm**4), ((0,), cm**2)])
+def test_apply_over_axes(axes, expected_units):
+    # the user-supplied function must be trusted to treat units
+    # sensibly (mainly that it doesn't give a mix of units across
+    # the resulting array), but we can check that units are
+    # propagated correctly for well-behaved functions.
     a = np.eye(3) * cm
-    ret = np.apply_over_axes(lambda x, axis: x * cm, a, (0, 1))
-    assert type(ret) is unyt_array
-    assert ret.units == cm**3
+    ret = np.apply_over_axes(lambda x, axis: x[axis] ** 2, a, axes)
+    assert isinstance(ret, unyt_array)  # could be unyt_quantity
+    assert ret.units == expected_units
 
 
 def test_array_equal():
@@ -1653,10 +1841,25 @@ def test_unique_values():
     assert_array_equal_units(res, values)
 
 
-def test_take():
+@pytest.mark.parametrize("indices", [[0, 1], 0])
+def test_take(indices):
     a = [1, 2, 3] * cm
-    res = np.take(a, [0, 1])
-    assert type(res) is unyt_array
+    res = np.take(a, indices)
+    if res.ndim == 0:
+        assert type(res) is unyt_quantity
+    else:
+        assert type(res) is unyt_array
+    assert res.units == cm
+
+
+@pytest.mark.parametrize("indices", [[0, 1], 0])
+def test_ndarray_take(indices):
+    a = [1, 2, 3] * cm
+    res = a.take(indices)
+    if res.ndim == 0:
+        assert type(res) is unyt_quantity
+    else:
+        assert type(res) is unyt_array
     assert res.units == cm
 
 
@@ -2089,19 +2292,19 @@ def test_all_arguments_reexposed(self, target, helper):
                         f"expected {kt}, got {kh}"
                     )
                 elif kt is KEYWORD_ONLY:
-                    assert (
-                        have_kwargs_helper
-                    ), f"argument {nt!r} is not re-exposed as keyword"
+                    assert have_kwargs_helper, (
+                        f"argument {nt!r} is not re-exposed as keyword"
+                    )
                 elif kt is POSITIONAL_OR_KEYWORD:
-                    assert (
-                        have_args_helper and have_kwargs_helper
-                    ), f"argument {nt!r} is not re-exposed as positional-or-keyword"
+                    assert have_args_helper and have_kwargs_helper, (
+                        f"argument {nt!r} is not re-exposed as positional-or-keyword"
+                    )
             elif kt is VAR_POSITIONAL:
                 assert have_args_helper, "helper is missing a catch-all *args argument"
             elif kt is VAR_KEYWORD:
-                assert (
-                    have_kwargs_helper
-                ), "helper is missing a catch-all **kwargs argument"
+                assert have_kwargs_helper, (
+                    "helper is missing a catch-all **kwargs argument"
+                )
 
     def test_known_arguments(self, target, helper):
         # validate that all exposed arguments map to something in the target
diff --git a/unyt/tests/test_unyt_array.py b/unyt/tests/test_unyt_array.py
index de0d1a3c..6841c0be 100644
--- a/unyt/tests/test_unyt_array.py
+++ b/unyt/tests/test_unyt_array.py
@@ -528,6 +528,44 @@ def test_power():
     assert_array_equal_units(cm_arr**0, unyt_quantity(1.0, "dimensionless"))
 
 
+@pytest.mark.parametrize(
+    "p",
+    [
+        pytest.param(unyt_array(0), id="scalar power 0"),
+        pytest.param(unyt_array(1), id="scalar power 1"),
+        pytest.param(unyt_array([0, 0, 0, 0]), id="uniform power 0"),
+        pytest.param(unyt_array([1, 1, 1, 1]), id="uniform power 1"),
+        pytest.param(unyt_array([0, 1, 2, 3]), id="non-uniform power"),
+    ],
+)
+def test_dimensionless_array_power(p):
+    # see https://github.com/yt-project/unyt/issues/522
+    a = unyt_array([1, 2, 3, 4])
+    expected = a.value**p.value
+    assert_array_equal(expected, (a**p).value)
+
+
+@pytest.mark.parametrize(
+    "a, p",
+    [
+        pytest.param(
+            unyt_array([0, 1, 2, 3], "cm"),
+            unyt_array([0, 1, 2, 3]),
+            id="non-uniform power with arr non dimensionless",
+        ),
+        pytest.param(
+            unyt_array([0, 1, 2, 3]),
+            unyt_array([0, 0, 0, 0], "cm"),
+            id="non dimensionless power",
+        ),
+    ],
+)
+def test_invalid_array_power(a, p):
+    # see https://github.com/yt-project/unyt/issues/522
+    with pytest.raises(UnitOperationError):
+        a**p
+
+
 def test_comparisons():
     """
     Test numpy ufunc comparison operators for unit consistency.