tensor_ops.hpp

Go to the documentation of this file.

 
#ifndef NUMERICAL_TENSOR_OPS_HPP
#define NUMERICAL_TENSOR_OPS_HPP
 
#include <algorithm>
#include <array>
#include <bit>
#include <cmath>
#include <complex>
#include <concepts>
#include <cstddef>
#include <cstring>
#include <functional>
#include <limits>
#include <string>
#include <type_traits>
#include <utility>
#include <vector>
 
#include "simd.hpp"
#include "tensor.hpp"
 
namespace numerical {
 
  namespace detail {
 
    [[nodiscard]] inline std::string shape_text(const std::vector<std::size_t>& shape)
    {
      std::string out {"("};
      for (std::size_t i {0}; i < shape.size(); ++i) {
        out += (i == 0 ? "" : ", ") + std::to_string(shape[i]);
      }
      return out + (shape.size() == 1 ? ",)" : ")");
    }
 
    [[nodiscard]] inline std::vector<std::size_t> broadcast_strides(const std::vector<std::size_t>& shape,
                                                                    std::size_t                     rank)
    {
      std::vector<std::size_t> strides(rank, 0);
      std::size_t              stride {1};
      for (std::size_t i {shape.size()}; i-- > 0;) {
        strides[rank - shape.size() + i] = shape[i] == 1 ? 0 : stride;
        stride                          *= shape[i];
      }
      return strides;
    }
 
    struct axis_walk
    {
      std::size_t outer;  
      std::size_t extent; 
      std::size_t inner;  
    };
 
    [[nodiscard]] inline axis_walk walk_along(const std::vector<std::size_t>& shape,
                                              std::size_t                     axis)
    {
      if (axis >= shape.size()) {
        throw dimension_error("axis " + std::to_string(axis) + " is out of bounds for a tensor of rank "
                              + std::to_string(shape.size()));
      }
      const auto before {shape.begin() + static_cast<std::ptrdiff_t>(axis)};
      axis_walk  walk   {.outer = 1, .extent = shape[axis], .inner = 1};
      // The result's elements, the axis taken out: none (and no line walked) where a kept extent is 0, whatever
      // the others; else their product, checked, which the outer and inner parts then divide.
      if (std::find(shape.begin(), before, std::size_t {0}) != before || std::find(before + 1, shape.end(), std::size_t {0}) != shape.end()) {
        walk.outer = 0;
        walk.inner = 0;
        return walk;
      }
      walk.outer = detail::element_count(shape.begin(), before);
      walk.inner = detail::element_count(before + 1, shape.end());
      static_cast<void>(detail::checked_mul(walk.outer, walk.inner));
      return walk;
    }
 
    template <std::floating_point R>
    [[nodiscard]] R numpy_muladd(R    x,
                                 R    y,
                                 R    z,
                                 bool fused)
    {
      if (fused) {
        return std::fma(x, y, z);
      }
      const R product {x * y};
      return product + z;
    }
 
    template <split_complex C>
    [[nodiscard]] C numpy_multiply(const C& a,
                                   const C& b,
                                   bool     fused)
    {
      using R = typename C::value_type;
      const R ii {a.imag() * b.imag()};
      const R ir {a.imag() * b.real()};
      return C {numpy_muladd(a.real(), b.real(), -ii, fused), numpy_muladd(a.real(), b.imag(), ir, fused)};
    }
 
    template <split_complex C>
    [[nodiscard]] C numpy_divide(const C& a,
                                 const C& b,
                                 bool     fused)
    {
      using R = typename C::value_type;
      const R real_part      {std::abs(b.real())};
      const R imaginary_part {std::abs(b.imag())};
      if (real_part >= imaginary_part) {
        if (real_part == R {0}) { // the imaginary part's magnitude no larger: a zero too
          return C {a.real() / real_part, a.imag() / imaginary_part};
        }
        const R ratio {b.imag() / b.real()};
        const R scale {R {1} / numpy_muladd(b.imag(), ratio, b.real(), fused)};
        return C {numpy_muladd(a.imag(), ratio, a.real(), fused) * scale, numpy_muladd(-a.real(), ratio, a.imag(), fused) * scale};
      }
      const R ratio {b.real() / b.imag()};
      const R scale {R {1} / numpy_muladd(b.real(), ratio, b.imag(), fused)};
      return C {numpy_muladd(a.real(), ratio, a.imag(), fused) * scale, numpy_muladd(a.imag(), ratio, -a.real(), fused) * scale};
    }
 
    [[nodiscard]] inline bool numpy_fuses_products() noexcept
    {
#if defined(__aarch64__) || defined(_M_ARM64)
      return true;
#elif defined(NUMERICAL_X86_COMPILED)
      return cpu_features().avx2_fma;
#else
      return false;
#endif
    }
 
    [[nodiscard]] inline bool numpy_fuses_scalar_products() noexcept
    {
#if defined(__APPLE__) || defined(_WIN32)
      return numpy_fuses_products();
#else
      return false;
#endif
    }
 
    [[nodiscard]] inline bool numpy_fuses_quotients() noexcept
    {
#if defined(__APPLE__) && defined(__aarch64__)
      return true;
#else
      return false;
#endif
    }
 
    template <typename T>
    [[nodiscard]] bool numpy_loop_is_scalar(const tensor<T>& a,
                                            const tensor<T>& b) noexcept
    {
      const auto        ones {[](const tensor<T>& t) { return std::ranges::all_of(t.shape(), [](std::size_t e) { return e == 1; }); }};
      const std::size_t rank {std::max(a.rank(), b.rank())};
      const auto        broadcast {[rank](const tensor<T>& t) { return t.rank() != 0 && t.rank() != rank; }};
      if (!ones(a) || !ones(b)) {
        return false;
      }
      return broadcast(a) || broadcast(b);
    }
 
    template <std::floating_point T>
    [[nodiscard]] T divided_by_count(T           value,
                                     std::size_t n)
    {
      using wide = std::common_type_t<T, double>;
      return static_cast<T>(static_cast<wide>(value) / static_cast<wide>(n));
    }
 
    template <split_complex T>
    [[nodiscard]] T divided_by_count(T           value,
                                     std::size_t n)
    {
      using R = typename T::value_type;
      const std::complex<double> wide {static_cast<double>(value.real()), static_cast<double>(value.imag())};
      const std::complex<double> mean {numpy_divide(wide, std::complex<double> {static_cast<double>(n), 0.0}, numpy_fuses_quotients())};
      return T {static_cast<R>(mean.real()), static_cast<R>(mean.imag())};
    }
 
    template <typename T>
    [[nodiscard]] T divided_by_count(T           value,
                                     std::size_t n)
    {
      if (n == 0) {
        throw dimension_error("a mean over no element has no value in an integer type");
      }
      return value / static_cast<T>(n);
    }
 
    template <std::floating_point T,
              typename Order>
    [[nodiscard]] T settled_zero(T            value,
                                 const T*     first,
                                 std::size_t  count,
                                 std::size_t  stride,
                                 const Order& larger)
    {
      if (value != T {0}) {
        return value;
      }
      const bool maximum {larger(T {1}, T {0})};
      for (std::size_t i {0}; i < count; ++i) {
        const T x {first[i * stride]};
        if (x == T {0} && std::signbit(x) != maximum) {
          return x;
        }
      }
      return value;
    }
 
    template <typename T,
              typename Order>
    [[nodiscard]] T settled_zero(T value,
                                 const T* /*first*/,
                                 std::size_t /*count*/,
                                 std::size_t /*stride*/,
                                 const Order& /*larger*/)
    {
      return value;
    }
 
    [[nodiscard]] inline std::vector<std::size_t> without_axis(std::vector<std::size_t> shape,
                                                               std::size_t              axis)
    {
      shape.erase(shape.begin() + static_cast<std::ptrdiff_t>(axis));
      return shape;
    }
 
    template <typename T,
              typename Reduce>
    [[nodiscard]] tensor<T> reduce_along(const tensor<T>& t,
                                         std::size_t      axis,
                                         const Reduce&    reduce)
    {
      const axis_walk walk                           {walk_along(t.shape(), axis)};
      buffer<T>       out(walk.outer * walk.inner);
      for (std::size_t o {0}; o < walk.outer; ++o) {
        for (std::size_t i {0}; i < walk.inner; ++i) {
          out[(o * walk.inner) + i] = reduce(t.data().data() + (o * walk.extent * walk.inner) + i, walk.extent, walk.inner);
        }
      }
      return tensor<T> {without_axis(t.shape(), axis), std::move(out)};
    }
 
    template <typename T>
    [[nodiscard]] T square(const T& x)
    {
      return static_cast<T>(x * x);
    }
 
    template <typename U>
    [[nodiscard]] std::complex<U> square(const std::complex<U>& x)
    {
      return {(x.real() * x.real()) - (x.imag() * x.imag()), (x.real() * x.imag()) + (x.imag() * x.real())};
    }
 
    template <typename T>
    [[nodiscard]] bool is_nan(const T& /*x*/)
    {
      return false;
    }
 
    [[nodiscard]] inline bool is_nan(double x)
    {
      return std::isnan(x);
    }
 
#if defined(__GNUC__)
    template <typename V,
              typename T>
    [[nodiscard]] auto better(const V& x,
                              const V& best,
                              const std::greater<T>& /*larger*/)
    {
      return x > best;
    }
 
    template <typename V,
              typename T>
    [[nodiscard]] auto better(const V& x,
                              const V& best,
                              const std::less<T>& /*smaller*/)
    {
      return x < best;
    }
 
    template <typename V>
    [[nodiscard]] auto unordered(const V& a,
                                 const V& b)
    {
      return a != b;
    }
 
#endif
 
    template <typename T,
              typename Order>
    [[nodiscard]] T contiguous_extreme(const T*     first,
                                       std::size_t  count,
                                       const Order& larger)
    {
      T              result   {first[0]};
      bool           nan      {false};
      std::size_t    done     {0};
#if defined(__GNUC__)
      using lane = typename native_vector<T>::type;
      constexpr std::size_t width {native_vector_bytes / sizeof(T)};
      constexpr std::size_t lanes {4}; // independent vectors, so no comparison waits on the one before it
      if (count >= width * lanes) {
        using mask = decltype(lane {} > lane {}); // all ones in a lane where a comparison holds
        std::array<lane, lanes> best {};
        std::array<mask, lanes> nans {};
        for (std::size_t l {0}; l < lanes; ++l) {
          std::memcpy(&best[l], first + (l * width), sizeof(lane));
        }
        done = count - (count % (width * lanes));
        for (std::size_t i {0}; i < done; i += width * lanes) {
#  pragma GCC unroll 4 // the lanes in registers: GCC keeps a rolled array of vectors in memory
          for (std::size_t l {0}; l < lanes; ++l) {
            lane x;
            std::memcpy(&x, first + i + (l * width), sizeof(lane));
            // Blended through the comparison's all-ones / all-zeros mask: both compilers keep it in vector
            // registers (GCC's `?:` on vectors does not).
            const mask take {better(x, best[l], larger)};
            best[l]  = std::bit_cast<lane>((std::bit_cast<mask>(x) & take) | (std::bit_cast<mask>(best[l]) & ~take));
            nans[l] |= unordered(x, x);
          }
        }
        for (std::size_t l {0}; l < lanes; ++l) {
          for (std::size_t k {0}; k < width; ++k) {
            nan    = nan || nans[l][k] != 0;
            result = larger(best[l][k], result) ? best[l][k] : result;
          }
        }
      }
#endif
      std::array<T, 8>        best {};
      std::array<unsigned, 8> nans {};
      best.fill(result);
      for (std::size_t i {done}; i < count; ++i) {
        const T x {first[i]};
        nans[i % 8] |= static_cast<unsigned>(std::isnan(x));
        best[i % 8]  = larger(x, best[i % 8]) ? x : best[i % 8];
      }
      for (std::size_t l {0}; l < 8; ++l) {
        nan    = nan || nans[l] != 0U;
        result = larger(best[l], result) ? best[l] : result;
      }
      return nan ? std::numeric_limits<T>::quiet_NaN() : result;
    }
 
    template <typename T,
              typename Order>
    [[nodiscard]] T extreme(const T*     first,
                            std::size_t  count,
                            std::size_t  stride,
                            const Order& larger)
    {
      if (count == 0) {
        throw dimension_error("a zero-size reduction to a maximum or a minimum has no identity");
      }
      if constexpr (std::is_floating_point_v<T>) {
        if (stride == 1) {
          // A contiguous line: the comparison without a branch (so it vectorises) and the NaN tracked apart,
          // a NaN anywhere is the answer, as NumPy's.
          return settled_zero(contiguous_extreme(first, count, larger), first, count, 1, larger);
        }
      }
      T best {first[0]};
      for (std::size_t i {1}; i < count; ++i) {
        // A NaN wins, as in NumPy: every comparison with it is false, so a NaN candidate is taken apart, and
        // once taken nothing compares larger than it.
        const T& candidate {first[i * stride]};
        if (is_nan(candidate) || larger(candidate, best)) {
          best = candidate;
        }
      }
      return settled_zero(best, first, count, stride, larger);
    }
 
    template <typename T,
              typename Order>
    [[nodiscard]] tensor<T> extremes_along(const tensor<T>& t,
                                           std::size_t      axis,
                                           const Order&     larger)
    {
      const axis_walk walk {walk_along(t.shape(), axis)};
      // An empty axis has no maximum even when the result is empty too (NumPy refuses both).
      if (walk.extent == 0) {
        throw dimension_error("a zero-size reduction to a maximum or a minimum has no identity");
      }
      if (!std::is_floating_point_v<T> || walk.inner == 1) {
        return reduce_along(t, axis, [&larger](const T* first, std::size_t count, std::size_t stride) {
                              return extreme(first, count, stride, larger);
                            });
      }
      buffer<T>             out(walk.outer * walk.inner);
      std::vector<unsigned> nan(walk.inner);
      for (std::size_t o {0}; o < walk.outer; ++o) {
        const T* const slab {t.data().data() + (o * walk.extent * walk.inner)};
        T* const       best {out.data() + (o * walk.inner)};
        std::copy(slab, slab + walk.inner, best);
        std::fill(nan.begin(), nan.end(), 0U);
        for (std::size_t e {0}; e < walk.extent; ++e) {
          const T* const row {slab + (e * walk.inner)};
          for (std::size_t i {0}; i < walk.inner; ++i) {
            const T x {row[i]};
            nan[i]  |= static_cast<unsigned>(std::isnan(x));
            best[i]  = larger(x, best[i]) ? x : best[i];
          }
        }
        for (std::size_t i {0}; i < walk.inner; ++i) {
          best[i] = settled_zero(best[i], slab + i, walk.extent, walk.inner, larger);
          if (nan[i] != 0U) {
            best[i] = std::numeric_limits<T>::quiet_NaN();
          }
        }
      }
      return tensor<T> {without_axis(t.shape(), axis), std::move(out)};
    }
  } // namespace detail
 
  [[nodiscard]] inline std::vector<std::size_t> broadcast_shape(const std::vector<std::size_t>& a,
                                                                const std::vector<std::size_t>& b)
  {
    const std::size_t        rank {std::max(a.size(), b.size())};
    std::vector<std::size_t> out(rank);
    for (std::size_t i {0}; i < rank; ++i) {
      const std::size_t x {i < rank - a.size() ? 1 : a[i - (rank - a.size())]};
      const std::size_t y {i < rank - b.size() ? 1 : b[i - (rank - b.size())]};
      if (x != y && x != 1 && y != 1) {
        throw dimension_error("operands could not be broadcast together with shapes " + detail::shape_text(a) + " "
                              + detail::shape_text(b));
      }
      out[i] = x == 1 ? y : x;
    }
    return out;
  }
 
  template <typename T,
            typename Op>
  [[nodiscard]] tensor<T> broadcast_apply(const tensor<T>& a,
                                          const tensor<T>& b,
                                          const Op&        op)
  {
    if (a.shape() == b.shape()) {
      buffer<T>         out(a.size());
      const std::size_t chunks {(a.size() + detail::elementwise_chunk - 1) / detail::elementwise_chunk};
      detail::parallel_for(chunks, detail::elementwise_chunk, [&](std::size_t c) {
                             const std::size_t first {c * detail::elementwise_chunk};
                             const std::size_t last {std::min(first + detail::elementwise_chunk, a.size())};
                             std::transform(a.data().begin() + static_cast<std::ptrdiff_t>(first), a.data().begin() + static_cast<std::ptrdiff_t>(last),
                                            b.data().begin() + static_cast<std::ptrdiff_t>(first), out.begin() + static_cast<std::ptrdiff_t>(first), op);
                           });
      return tensor<T> {a.shape(), std::move(out)};
    }
    std::vector<std::size_t>       shape {broadcast_shape(a.shape(), b.shape())};
    const std::size_t              rank  {shape.size()};
    const std::vector<std::size_t> sa    {detail::broadcast_strides(a.shape(), rank)};
    const std::vector<std::size_t> sb    {detail::broadcast_strides(b.shape(), rank)};
    const std::size_t              count {detail::element_count(shape.begin(), shape.end())};
    buffer<T>                      out(count); // every element written below
    if (count == 0) {
      return tensor<T> {std::move(shape), std::move(out)};
    }
    // The innermost axis runs as one loop, each operand stepping 1 or 0 (a broadcast one) along it, so a
    // scalar or a row against a matrix is a plain vector loop; the odometer turns only between rows. The rows go
    // by tasks, each task's odometer set from its first row.
    const std::size_t inner          {shape[rank - 1]};
    const std::size_t step_a         {sa[rank - 1]};
    const std::size_t step_b         {sb[rank - 1]};
    const std::size_t rows           {count / inner};
    const std::size_t rows_per_task  {std::max<std::size_t>(1, detail::elementwise_chunk / inner)};
    detail::parallel_for((rows + rows_per_task - 1) / rows_per_task, rows_per_task * inner, [&](std::size_t task) {
                           const std::size_t        first {task * rows_per_task};
                           const std::size_t        last {std::min(first + rows_per_task, rows)};
                           std::vector<std::size_t> index(rank, 0);
                           std::size_t              ia {0};
                           std::size_t              ib {0};
                           for (std::size_t axis {rank - 1}, rest {first}; axis-- > 0;) { // the first row's position
                             index[axis]  = rest % shape[axis];
                             rest        /= shape[axis];
                             ia          += index[axis] * sa[axis];
                             ib          += index[axis] * sb[axis];
                           }
                           for (std::size_t row {first}; row < last; ++row) {
                             const T* const pa {a.data().data() + ia};
                             const T* const pb {b.data().data() + ib};
                             T* const       po {out.data() + (row * inner)};
                             if (step_a == 1 && step_b == 1) {
                               for (std::size_t i {0}; i < inner; ++i) {
                                 po[i] = op(pa[i], pb[i]);
                               }
                             }
                             else if (step_a == 1) {
                               const T y {pb[0]};
                               for (std::size_t i {0}; i < inner; ++i) {
                                 po[i] = op(pa[i], y);
                               }
                             }
                             else if (step_b == 1) {
                               const T x {pa[0]};
                               for (std::size_t i {0}; i < inner; ++i) {
                                 po[i] = op(x, pb[i]);
                               }
                             }
                             else {
                               std::fill(po, po + inner, op(pa[0], pb[0]));
                             }
                             for (std::size_t axis {rank - 1}; axis-- > 0;) { // odometer over the outer axes: carry on wrap
                               ++index[axis];
                               ia += sa[axis];
                               ib += sb[axis];
                               if (index[axis] < shape[axis]) {
                                 break;
                               }
                               ia          -= sa[axis] * shape[axis];
                               ib          -= sb[axis] * shape[axis];
                               index[axis]  = 0;
                             }
                           }
                         });
    return tensor<T> {std::move(shape), std::move(out)};
  }
 
  template <detail::split_complex C>
  [[nodiscard]] tensor<C> numpy_multiply(const tensor<C>& a,
                                         const tensor<C>& b)
  {
    const bool fused {detail::numpy_loop_is_scalar(a, b) ? detail::numpy_fuses_scalar_products() : detail::numpy_fuses_products()};
    return broadcast_apply(a, b, [fused](const C& x, const C& y) { return detail::numpy_multiply(x, y, fused); });
  }
 
  template <detail::split_complex C>
  [[nodiscard]] tensor<C> numpy_divide(const tensor<C>& a,
                                       const tensor<C>& b)
  {
    const bool fused {detail::numpy_fuses_quotients()};
    return broadcast_apply(a, b, [fused](const C& x, const C& y) { return detail::numpy_divide(x, y, fused); });
  }
 
  template <typename T>
  [[nodiscard]] tensor<T> operator+(const tensor<T>& a,
                                    const tensor<T>& b)
  {
    return broadcast_apply(a, b, std::plus<T> {});
  }
 
  template <typename T>
  [[nodiscard]] tensor<T> operator-(const tensor<T>& a,
                                    const tensor<T>& b)
  {
    return broadcast_apply(a, b, std::minus<T> {});
  }
 
  template <typename T>
  [[nodiscard]] tensor<T> operator*(const tensor<T>& a,
                                    const tensor<T>& b)
  {
    return broadcast_apply(a, b, std::multiplies<T> {});
  }
 
  template <typename T>
  [[nodiscard]] tensor<T> operator/(const tensor<T>& a,
                                    const tensor<T>& b)
  {
    return broadcast_apply(a, b, std::divides<T> {});
  }
 
  template <typename T>
  [[nodiscard]] tensor<T> operator-(const tensor<T>& a)
  {
    buffer<T> out(a.size());
    std::transform(a.data().begin(), a.data().end(), out.begin(), std::negate<T> {});
    return tensor<T> {a.shape(), std::move(out)};
  }
 
  template <typename T,
            typename F>
  [[nodiscard]] auto map(const tensor<T>& t,
                         const F&         f) -> tensor<std::decay_t<decltype(f(std::declval<const T&>()))>>
  {
    using R = std::decay_t<decltype(f(std::declval<const T&>()))>;
    buffer<R> out(t.size());
    std::transform(t.data().begin(), t.data().end(), out.begin(), f);
    return tensor<R> {t.shape(), std::move(out)};
  }
 
  template <typename T>
  [[nodiscard]] tensor<T> pow(const tensor<T>& a,
                              const tensor<T>& b)
  {
    if constexpr (!std::is_integral_v<T>) {
      if (b.size() == 1 && b.shape().size() <= a.shape().size()) {
        const T y {b.data()[0]};
        if (y == T(2)) {
          return map(a, [](const T& x) { return detail::square(x); });
        }
        if (y == T(0.5)) {
          return map(a, [](const T& x) { return static_cast<T>(std::sqrt(x)); });
        }
        if (y == T(-1)) {
          return map(a, [](const T& x) { return static_cast<T>(T(1) / x); });
        }
      }
    }
    return broadcast_apply(a, b, [](const T& x, const T& y) { return static_cast<T>(std::pow(x, y)); });
  }
 
  template <typename T>
  [[nodiscard]] T sum(const tensor<T>& t)
  {
    return detail::pairwise_sum(t.data().data(), t.size(), 1);
  }
 
  template <typename T>
  [[nodiscard]] tensor<T> sum(const tensor<T>& t,
                              std::size_t      axis)
  {
    const detail::axis_walk walk {detail::walk_along(t.shape(), axis)};
    buffer<T>               out(walk.outer * walk.inner);
    detail::pairwise_sums(t.data().data(), walk.outer, walk.extent, walk.inner, out.data());
    return tensor<T> {detail::without_axis(t.shape(), axis), std::move(out)};
  }
 
  template <typename T>
  [[nodiscard]] T mean(const tensor<T>& t)
  {
    return detail::divided_by_count(sum(t), t.size());
  }
 
  template <typename T>
  [[nodiscard]] tensor<T> mean(const tensor<T>& t,
                               std::size_t      axis)
  {
    const detail::axis_walk walk  {detail::walk_along(t.shape(), axis)};
    buffer<T>               out(walk.outer * walk.inner);
    detail::pairwise_sums(t.data().data(), walk.outer, walk.extent, walk.inner, out.data());
    for (T& x : out) {
      x = detail::divided_by_count(x, walk.extent);
    }
    return tensor<T> {detail::without_axis(t.shape(), axis), std::move(out)};
  }
 
  template <std::totally_ordered T>
  [[nodiscard]] T max(const tensor<T>& t)
  {
    return detail::extreme(t.data().data(), t.size(), 1, std::greater<T> {});
  }
 
  template <std::totally_ordered T>
  [[nodiscard]] T min(const tensor<T>& t)
  {
    return detail::extreme(t.data().data(), t.size(), 1, std::less<T> {});
  }
 
  template <std::totally_ordered T>
  [[nodiscard]] tensor<T> max(const tensor<T>& t,
                              std::size_t      axis)
  {
    return detail::extremes_along(t, axis, std::greater<T> {});
  }
 
  template <std::totally_ordered T>
  [[nodiscard]] tensor<T> min(const tensor<T>& t,
                              std::size_t      axis)
  {
    return detail::extremes_along(t, axis, std::less<T> {});
  }
 
  template <typename T>
  [[nodiscard]] tensor<T> reshape(const tensor<T>&         t,
                                  std::vector<std::size_t> shape)
  {
    std::size_t count {1};
    for (const std::size_t extent : shape) {
      count = detail::checked_mul(count, extent);
    }
    if (count != t.size()) {
      throw dimension_error("cannot reshape a tensor of size " + std::to_string(t.size()) + " into shape "
                            + detail::shape_text(shape));
    }
    return tensor<T> {std::move(shape), t.data()};
  }
 
  template <typename T>
  [[nodiscard]] tensor<T> full(std::vector<std::size_t> shape,
                               const T&                 value)
  {
    std::size_t count {1};
    for (const std::size_t extent : shape) {
      count = detail::checked_mul(count, extent);
    }
    return tensor<T> {std::move(shape), buffer<T>(count, value)};
  }
 
  template <typename T>
  [[nodiscard]] tensor<T> zeros(std::vector<std::size_t> shape)
  {
    return full<T>(std::move(shape), T {});
  }
 
  template <typename T>
  [[nodiscard]] tensor<T> ones(std::vector<std::size_t> shape)
  {
    return full<T>(std::move(shape), static_cast<T>(1));
  }
 
  template <typename T>
  [[nodiscard]] tensor<T> eye(std::size_t n)
  {
    buffer<T> out(detail::checked_mul(n, n), T {});
    for (std::size_t i {0}; i < n; ++i) {
      out[(i * n) + i] = static_cast<T>(1);
    }
    return tensor<T> {{n, n}, std::move(out)};
  }
 
  [[nodiscard]] inline tensor<double> arange(double start,
                                             double stop,
                                             double step)
  {
    if (step == 0.0 || !std::isfinite(step) || !std::isfinite(start) || !std::isfinite(stop)) {
      throw dimension_error("arange needs finite bounds and a finite, non-zero step");
    }
    const double        length            {std::ceil((stop - start) / step)};
    if (!(length < 0x1p64)) { // past std::size_t, or an infinite span: no conversion holds it
      throw dimension_error("arange: the length does not fit in std::size_t");
    }
    const std::size_t   count             {length > 0.0 ? static_cast<std::size_t>(length) : 0};
    buffer<double>      out(count);
    const double        delta             {(start + step) - start};
    for (std::size_t i {0}; i < count; ++i) {
      out[i] = std::fma(static_cast<double>(i), delta, start);
    }
    return tensor<double> {{count}, std::move(out)};
  }
 
  [[nodiscard]] inline tensor<double> linspace(double      start,
                                               double      stop,
                                               std::size_t count)
  {
    buffer<double>      out(count);
    const double        step              {count > 1 ? (stop - start) / static_cast<double>(count - 1) : 0.0};
    for (std::size_t i {0}; i < count; ++i) {
      // The product rounded before the sum, as NumPy's linspace computes it in two operations. A compiler
      // may fuse `i * step` into the addition after it (GCC does, across statements, wherever there is a
      // fused multiply-add); a product computed by fma, `i * step - 0`, is exactly the rounded one, signed
      // zero included, and no multiplication is left to fuse.
      const double offset {std::fma(static_cast<double>(i), step, -0.0)};
      out[i] = start + offset;
    }
    if (count > 1) {
      out[count - 1] = stop;
    }
    return tensor<double> {{count}, std::move(out)};
  }
} // namespace numerical
 
#endif // NUMERICAL_TENSOR_OPS_HPP