tensor.hpp

Go to the documentation of this file.

 
#ifndef NUMERICAL_TENSOR_HPP
#define NUMERICAL_TENSOR_HPP
 
#include <algorithm>
#include <array>
#include <cctype>
#include <cmath>
#include <complex> // so std::abs resolves for a complex scalar, regardless of include order
#include <cstddef>
#include <deque>
#include <initializer_list>
#include <map>
#include <optional>
#include <stdexcept>
#include <string>
#include <string_view>
#include <utility>
#include <vector>
 
#include "gemm.hpp"
#include "matrix.hpp"
#include "parallel.hpp"
#include "vector.hpp"
 
namespace numerical {
 
  class einsum_error : public std::runtime_error
  {
  public:
 
    explicit einsum_error(const std::string& message)
      : std::runtime_error(message)
    {}
  };
 
  template <typename T>
  class tensor
  {
  public:
 
    tensor(std::vector<std::size_t> shape,
           buffer<T>                elements)
      : shape_(std::move(shape)), data_(std::move(elements))
    {
      finalize();
    }
 
    tensor(std::vector<std::size_t> shape,
           const std::vector<T>&    elements)
      : tensor(std::move(shape),
               buffer<T>(elements.begin(), elements.end()))
    {}
 
    tensor(std::vector<std::size_t> shape,
           std::initializer_list<T> elements)
      : tensor(std::move(shape),
               buffer<T>(elements))
    {}
 
    explicit tensor(const vector<T>& v)
      : shape_ {v.size()}
    {
      data_.reserve(v.size());
      for (std::size_t i {0}; i < v.size(); ++i) {
        data_.push_back(v[i]);
      }
      finalize();
    }
 
    explicit tensor(matrix<T>&& m)
      : shape_ {m.rows(), m.cols()}, data_ {std::move(m).elements()}
    {
      finalize();
    }
 
    explicit tensor(const matrix<T>& m)
      : shape_ {m.rows(), m.cols()}, data_(m.begin(), m.end())
    {
      finalize();
    }
 
    [[nodiscard]] std::size_t rank() const noexcept
    {
      return shape_.size();
    }
 
    [[nodiscard]] std::size_t size() const noexcept
    {
      return data_.size();
    }
 
    [[nodiscard]] const std::vector<std::size_t>& shape() const noexcept
    {
      return shape_;
    }
 
    [[nodiscard]] const buffer<T>& data() const noexcept
    {
      return data_;
    }
 
    [[nodiscard]] const T& operator()(const std::vector<std::size_t>& index) const
    {
      std::size_t flat {0};
      for (std::size_t axis {0}; axis < index.size(); ++axis) {
        flat += strides_[axis] * index[axis];
      }
      return data_[flat];
    }
 
    [[nodiscard]] bool operator==(const tensor& other) const
    {
      return shape_ == other.shape_ && data_ == other.data_;
    }
 
  private:
 
    void finalize()
    {
      strides_.assign(shape_.size(), 1);
      std::size_t acc {1};
      for (std::size_t axis {shape_.size()}; axis-- > 0;) {
        strides_[axis] = acc;
        acc            = detail::checked_mul(acc, shape_[axis]); // guard the element-count product
      }
      if (data_.size() != acc) {
        throw dimension_error("tensor: data size does not match the shape");
      }
    }
 
    std::vector<std::size_t> shape_;   
    std::vector<std::size_t> strides_; 
    buffer<T>                data_;    
  };
 
  template <typename T>
  [[nodiscard]] bool approx_equal(const tensor<T>& a,
                                  const tensor<T>& b,
                                  double           tolerance = 1e-9)
  {
    if (a.shape() != b.shape()) {
      return false;
    }
    for (std::size_t i {0}; i < a.size(); ++i) {
      // Equal, or within the tolerance: a NaN is neither, so it is never close, and two equal
      // infinities are equal although their difference is NaN.
      if (!(a.data()[i] == b.data()[i]) && !(std::abs(a.data()[i] - b.data()[i]) <= tolerance)) {
        return false;
      }
    }
    return true;
  }
 
  namespace detail {
    template <typename T>
    [[nodiscard]] T pairwise_sum(const T*    first,
                                 std::size_t count,
                                 std::size_t stride)
    {
      constexpr std::size_t lanes {8};
      constexpr std::size_t block {128};
      if (count < lanes) {
        T sum {};
        for (std::size_t i {0}; i < count; ++i) {
          sum += first[i * stride];
        }
        return sum;
      }
      if (count <= block) {
        std::array<T, lanes> partial {};
        for (std::size_t j {0}; j < lanes; ++j) {
          partial[j] = first[j * stride];
        }
        std::size_t i {lanes};
        for (; i < count - (count % lanes); i += lanes) {
          for (std::size_t j {0}; j < lanes; ++j) {
            partial[j] += first[(i + j) * stride];
          }
        }
        T sum {((partial[0] + partial[1]) + (partial[2] + partial[3])) + ((partial[4] + partial[5]) + (partial[6] + partial[7]))};
        for (; i < count; ++i) {
          sum += first[i * stride];
        }
        return sum;
      }
      std::size_t half {count / 2};
      half -= half % lanes;
      return pairwise_sum(first, half, stride) + pairwise_sum(first + (half * stride), count - half, stride);
    }
 
    template <typename T>
    void pairwise_sum_rows(const T*    first,
                           std::size_t count,
                           std::size_t stride,
                           std::size_t width,
                           T*          out)
    {
      constexpr std::size_t lanes {8};
      constexpr std::size_t block {128};
      const auto            row = [&](std::size_t i) { return first + (i * stride); };
      if (count < lanes) {
        std::fill_n(out, width, T {});
        for (std::size_t i {0}; i < count; ++i) {
          for (std::size_t q {0}; q < width; ++q) {
            out[q] += row(i)[q];
          }
        }
        return;
      }
      if (count <= block) {
        std::vector<T> partial(lanes * width);
        for (std::size_t j {0}; j < lanes; ++j) {
          std::copy_n(row(j), width, partial.data() + (j * width));
        }
        std::size_t i {lanes};
        for (; i < count - (count % lanes); i += lanes) {
          for (std::size_t j {0}; j < lanes; ++j) {
            T* const       sums {partial.data() + (j * width)};
            const T* const next {row(i + j)};
            for (std::size_t q {0}; q < width; ++q) {
              sums[q] += next[q];
            }
          }
        }
        const T* const p {partial.data()};
        for (std::size_t q {0}; q < width; ++q) {
          out[q] = ((p[q] + p[width + q]) + (p[(2 * width) + q] + p[(3 * width) + q]))
                   + ((p[(4 * width) + q] + p[(5 * width) + q]) + (p[(6 * width) + q] + p[(7 * width) + q]));
        }
        for (; i < count; ++i) {
          for (std::size_t q {0}; q < width; ++q) {
            out[q] += row(i)[q];
          }
        }
        return;
      }
      std::size_t half {count / 2};
      half -= half % lanes;
      std::vector<T> second(width);
      pairwise_sum_rows(first, half, stride, width, out);
      pairwise_sum_rows(row(half), count - half, stride, width, second.data());
      for (std::size_t q {0}; q < width; ++q) {
        out[q] = out[q] + second[q];
      }
    }
 
    inline constexpr std::size_t pairwise_columns {512};
 
    template <typename T>
    void pairwise_sums(const T*    data,
                       std::size_t outer,
                       std::size_t count,
                       std::size_t inner,
                       T*          out)
    {
      for (std::size_t o {0}; o < outer; ++o) {
        const T* const base {data + (o * count * inner)};
        if (inner == 1) {
          out[o] = pairwise_sum(base, count, 1);
          continue;
        }
        for (std::size_t q {0}; q < inner; q += pairwise_columns) {
          pairwise_sum_rows(base + q, count, inner, std::min(pairwise_columns, inner - q), out + (o * inner) + q);
        }
      }
    }
 
    [[nodiscard]] inline bool is_label(char ch)
    {
      return std::isalpha(static_cast<unsigned char>(ch)) != 0;
    }
 
    [[nodiscard]] inline std::vector<std::string> split_inputs(const std::string& lhs)
    {
      std::vector<std::string> groups;
      std::string              current;
      for (const char ch : lhs) {
        if (ch == ',') {
          groups.push_back(current);
          current.clear();
        }
        else if (is_label(ch)) {
          current.push_back(ch);
        }
        else {
          throw einsum_error("einsum: invalid character in input subscripts");
        }
      }
      groups.push_back(current);
      return groups;
    }
 
    template <typename T>
    [[nodiscard]] tensor<T> tile_gemm(const T*                 a,
                                      const T*                 b,
                                      std::size_t              rows,
                                      std::size_t              inner,
                                      std::size_t              cols,
                                      std::vector<std::size_t> out_shape)
    {
      buffer<T> c(rows * cols);          // the caller routes only an output of two elements or more here
      parallel_zero(c.data(), c.size()); // gemm adds to C
      if (inner != 0) {
        gemm(rows, cols, inner, a, inner, b, cols, c.data(), cols);
      }
      return tensor<T> {std::move(out_shape), std::move(c)};
    }
 
    inline constexpr std::size_t vector_chunk {128};
 
    template <typename T>
    [[nodiscard]] tensor<T> canonical_product(const T*                 a,
                                              const T*                 b,
                                              std::size_t              rows,
                                              std::size_t              inner,
                                              std::size_t              cols,
                                              std::vector<std::size_t> out_shape)
    {
      if (cols == 1 && rows > 1 && inner != 0) { // a matrix times a vector: each row's dot product, as "i,i->"
        buffer<T> y(rows);
        matrix_times_vector(a, rows, inner, b, y.data());
        return tensor<T> {std::move(out_shape), std::move(y)};
      }
      if (rows == 1 && cols > 1 && inner != 0) { // a vector times a matrix: its rows scaled and added in order
        buffer<T> y(cols, T {});
        parallel_for((cols + vector_chunk - 1) / vector_chunk, inner * vector_chunk, [&](std::size_t chunk) {
                       const std::size_t first {chunk * vector_chunk};
                       const std::size_t count {std::min(vector_chunk, cols - first)};
                       for (std::size_t k {0}; k < inner; ++k) {
                         add_product(y.data() + first, a[k], b + (k * cols) + first, count);
                       }
                     });
        return tensor<T> {std::move(out_shape), std::move(y)};
      }
      if (rows * cols != 1) { // the product, an empty one included; a single element is one dot product
        return tile_gemm<T>(a, b, rows, inner, cols, std::move(out_shape));
      }
      return tensor<T> {std::move(out_shape), buffer<T> {chunked_dot(a, b, inner)}};
    }
 
    inline constexpr std::size_t transposed_columns {8};
 
    template <typename T>
    [[nodiscard]] buffer<T> vector_times_transposed(const T*    a,
                                                    const T*    bt,
                                                    std::size_t ldbt,
                                                    std::size_t inner,
                                                    std::size_t cols)
    {
      buffer<T> y(cols, T {});
      parallel_for((cols + vector_chunk - 1) / vector_chunk, inner * vector_chunk, [&](std::size_t chunk) {
                     const std::size_t last {std::min((chunk + 1) * vector_chunk, cols)};
                     std::size_t j {chunk * vector_chunk};
                     for (; j + transposed_columns <= last; j += transposed_columns) { // a whole group: its sums in registers
                       std::array<T, transposed_columns> sums {};
                       for (std::size_t k {0}; k < inner; ++k) {
                         const T f {a[k]};
                         for (std::size_t c {0}; c < transposed_columns; ++c) {
                           add_one_product(sums[c], f, bt[((j + c) * ldbt) + k]);
                         }
                       }
                       std::copy(sums.begin(), sums.end(), y.data() + j);
                     }
                     for (; j < last; ++j) {
                       T sum {};
                       for (std::size_t k {0}; k < inner; ++k) {
                         add_one_product(sum, a[k], bt[(j * ldbt) + k]);
                       }
                       y[j] = sum;
                     }
                   });
      return y;
    }
 
#ifdef NUMERICAL_X86_COMPILED
    [[nodiscard]] inline buffer<double> vector_times_transposed(const double* a,
                                                                const double* bt,
                                                                std::size_t   ldbt,
                                                                std::size_t   inner,
                                                                std::size_t   cols)
    {
      const gemm_path path {gemm_path_current()};
      if (path != gemm_path::avx2 && path != gemm_path::avx512) {
        return vector_times_transposed<double>(a, bt, ldbt, inner, cols);
      }
      buffer<double> y(cols);
      parallel_for((cols + vector_chunk - 1) / vector_chunk, inner * vector_chunk, [&](std::size_t chunk) {
                     avx2_transposed_columns(a, bt, ldbt, inner, chunk * vector_chunk, std::min((chunk + 1) * vector_chunk, cols), y.data());
                   });
      return y;
    }
 
#endif
 
    template <typename T>
    [[nodiscard]] std::optional<tensor<T>> sum_over_run(const std::string&              in,
                                                        const std::string&              output,
                                                        const std::vector<std::size_t>& shape,
                                                        const T*                        data)
    {
      const std::size_t first {in.find_first_not_of(output)};
      if (first == std::string::npos) {
        return std::nullopt; // nothing summed: a permutation, or the operand itself
      }
      const std::size_t next {in.find_first_of(output, first)};
      const std::size_t past {next == std::string::npos ? in.size() : next};
      if (in.substr(0, first) + in.substr(past) != output) {
        return std::nullopt;
      }
      std::vector<std::size_t> out_shape;
      for (std::size_t k {0}; k < in.size(); ++k) {
        if (k < first || k >= past) {
          out_shape.push_back(shape[k]);
        }
      }
      // The result's elements: none where a kept extent is 0, whatever the others (each part then 0); else their
      // checked product, which bounds the parts, and the run's with them the operand's.
      const std::size_t kept  {element_count(out_shape.begin(), out_shape.end())};
      if (kept == 0) {
        return tensor<T> {std::move(out_shape), buffer<T> {}};
      }
      const auto        begin {shape.begin()};
      const std::size_t outer {element_count(begin, begin + static_cast<std::ptrdiff_t>(first))};
      const std::size_t run   {element_count(begin + static_cast<std::ptrdiff_t>(first), begin + static_cast<std::ptrdiff_t>(past))};
      const std::size_t inner {element_count(begin + static_cast<std::ptrdiff_t>(past), shape.end())};
      std::vector<T>    sums(kept);
      pairwise_sums(data, outer, run, inner, sums.data());
      return tensor<T> {std::move(out_shape), std::move(sums)};
    }
 
    template <typename T>
    [[nodiscard]] buffer<T> permuted(const buffer<T>&                   data,
                                     const std::string&                 from,
                                     const std::string&                 to,
                                     const std::map<char, std::size_t>& sizes)
    {
      buffer<T> out(data.size()); // every element written below
      if (data.empty()) {
        return out;
      }
      if (to.size() == 2) { // the one permutation of two axes, a transpose: by tiles, each read and written in the cache
        const std::size_t rows {sizes.at(from[0])};
        const std::size_t cols {sizes.at(from[1])};
        parallel_for((rows + transpose_tile - 1) / transpose_tile, transpose_tile * cols, [&](std::size_t block) {
                       const std::size_t r0 {block * transpose_tile};
                       const std::size_t r1 {std::min(r0 + transpose_tile, rows)};
                       for (std::size_t c0 {0}; c0 < cols; c0 += transpose_tile) {
                         const std::size_t c1 {std::min(c0 + transpose_tile, cols)};
                         for (std::size_t r {r0}; r < r1; ++r) {
                           for (std::size_t c {c0}; c < c1; ++c) {
                             out[(c * rows) + r] = data[(r * cols) + c];
                           }
                         }
                       }
                     });
        return out;
      }
      const std::size_t           rank {to.size()};
      std::map<char, std::size_t> source; // each label's stride in data
      std::size_t                 step {1};
      for (std::size_t axis {from.size()}; axis-- > 0;) {
        source[from[axis]]  = step;
        step               *= sizes.at(from[axis]);
      }
      std::vector<std::size_t> stride(rank);
      std::vector<std::size_t> extent(rank);
      for (std::size_t axis {0}; axis < rank; ++axis) {
        stride[axis] = source.at(to[axis]);
        extent[axis] = sizes.at(to[axis]);
      }
      const std::size_t        inner  {extent[rank - 1]};
      const std::size_t        sweep  {stride[rank - 1]};
      std::vector<std::size_t> index(rank - 1, 0);
      std::size_t              offset {0};
      for (std::size_t row {0}; row < data.size() / inner; ++row) {
        for (std::size_t i {0}; i < inner; ++i) {
          out[(row * inner) + i] = data[offset + (i * sweep)];
        }
        for (std::size_t axis {rank - 1}; axis-- > 0;) {
          offset += stride[axis];
          if (++index[axis] < extent[axis]) {
            break;
          }
          offset      -= stride[axis] * extent[axis];
          index[axis]  = 0;
        }
      }
      return out;
    }
 
    template <typename T>
    [[nodiscard]] std::optional<tensor<T>> contract_pair(const std::string&                 sa,
                                                         const tensor<T>&                   a,
                                                         const std::string&                 sb,
                                                         const tensor<T>&                   b,
                                                         const std::string&                 output,
                                                         const std::string&                 summed,
                                                         const std::map<char, std::size_t>& sizes)
    {
      const auto repeats {[](const std::string& labels) {
                            for (std::size_t i {0}; i < labels.size(); ++i) {
                              if (labels.find(labels[i], i + 1) != std::string::npos) {
                                return true;
                              }
                            }
                            return false;
                          }};
      if (repeats(sa) || repeats(sb)) {
        return std::nullopt;
      }
      std::string batch;
      std::string free_a;
      std::string free_b;
      for (const char label : output) {
        const bool in_a {sa.find(label) != std::string::npos};
        const bool in_b {sb.find(label) != std::string::npos};
        (in_a && in_b ? batch : (in_a ? free_a : free_b)).push_back(label);
      }
      if (batch + free_a + free_b != output && batch + free_b + free_a == output) {
        // The other operand first lays the result out in the output's order: the same terms in the same order
        // (each a product of the same two numbers), with no permutation of the result.
        return contract_pair(sb, b, sa, a, output, summed, sizes);
      }
      std::size_t inner {1};
      for (const char label : summed) {
        if (sa.find(label) == std::string::npos || sb.find(label) == std::string::npos) {
          return std::nullopt;
        }
        inner *= sizes.at(label);
      }
      const auto extent_of {[&sizes](const std::string& labels) {
                              std::size_t product {1};
                              for (const char label : labels) {
                                product = checked_mul(product, sizes.at(label));
                              }
                              return product;
                            }};
      const std::size_t    batches                                            {extent_of(batch)};
      const std::size_t    rows                                               {extent_of(free_a)};
      const std::size_t    cols                                               {extent_of(free_b)};
      const std::string    left_order                                         {batch + free_a + summed};
      const std::string    right_order                                        {batch + summed + free_b};
      const buffer<T>      left                                               {sa == left_order ? buffer<T> {} : permuted(a.data(), sa, left_order, sizes)};
      const buffer<T>      right                                              {sb == right_order ? buffer<T> {} : permuted(b.data(), sb, right_order, sizes)};
      const T* const       left_data                                          {sa == left_order ? a.data().data() : left.data()};
      const T* const       right_data                                         {sb == right_order ? b.data().data() : right.data()};
      buffer<T>            c(checked_mul(batches, checked_mul(rows, cols)));
      if (rows == 1 && cols == 1) { // a dot product per batch, as `chunked_dot` sums each chunk
        parallel_for(batches, inner, [&](std::size_t p) {
                       c[p] = dot_product(left_data + (p * inner), right_data + (p * inner), inner);
                     });
      }
      else {
        parallel_zero(c.data(), c.size()); // gemm adds to C
        parallel_for(batches, rows * cols * inner / gemm_kernel_speedup, [&](std::size_t p) {
                       gemm(rows, cols, inner, left_data + (p * rows * inner), inner, right_data + (p * inner * cols), cols,
                            c.data() + (p * rows * cols), cols);
                     });
      }
      std::vector<std::size_t> shape;
      for (const char label : output) {
        shape.push_back(sizes.at(label));
      }
      const std::string order {batch + free_a + free_b};
      return tensor<T> {std::move(shape), order == output ? std::move(c) : permuted(c, order, output, sizes)};
    }
 
    inline constexpr std::size_t exhaustive_path_operands {4};
 
    struct contraction_step
    {
      std::size_t first;  
      std::size_t second; 
    };
 
    [[nodiscard]] inline std::string pair_keeps(const std::vector<std::string>& labels,
                                                std::size_t                     i,
                                                std::size_t                     j,
                                                const std::string&              output)
    {
      std::string elsewhere {output};
      for (std::size_t k {0}; k < labels.size(); ++k) {
        elsewhere += k == i || k == j ? std::string {} : labels[k];
      }
      std::string keep;
      for (const char label : labels[i] + labels[j]) {
        if (elsewhere.find(label) != std::string::npos && keep.find(label) == std::string::npos) {
          keep.push_back(label);
        }
      }
      return keep;
    }
 
    [[nodiscard]] inline double extent_product(const std::string&                 labels,
                                               const std::map<char, std::size_t>& sizes)
    {
      double      product {1.0};
      std::string seen;
      for (const char label : labels) {
        if (seen.find(label) == std::string::npos) {
          seen.push_back(label);
          product *= static_cast<double>(sizes.at(label));
        }
      }
      return product;
    }
 
    [[nodiscard]] inline std::vector<std::string> after_step(const std::vector<std::string>& labels,
                                                             contraction_step                step,
                                                             const std::string&              output)
    {
      std::vector<std::string> next;
      for (std::size_t k {0}; k < labels.size(); ++k) {
        if (k != step.first && k != step.second) {
          next.push_back(labels[k]);
        }
      }
      next.push_back(pair_keeps(labels, step.first, step.second, output));
      return next;
    }
 
    inline void cheapest_path(const std::vector<std::string>&    labels,
                              const std::string&                 output,
                              const std::map<char, std::size_t>& sizes,
                              std::vector<contraction_step>&     path,
                              double&                            cost,
                              double&                            largest)
    {
      if (labels.size() < 2) {
        path.clear();
        cost    = 0.0;
        largest = 0.0;
        return;
      }
      bool found {false};
      for (std::size_t i {0}; i < labels.size(); ++i) {
        for (std::size_t j {i + 1}; j < labels.size(); ++j) {
          const contraction_step        step         {.first = i, .second = j};
          std::vector<contraction_step> rest;
          double                        rest_cost    {0.0};
          double                        rest_largest {0.0};
          cheapest_path(after_step(labels, step, output), output, sizes, rest, rest_cost, rest_largest);
          const double here                          {extent_product(labels[i] + labels[j], sizes)};
          const double size                          {extent_product(pair_keeps(labels, i, j, output), sizes)};
          const double total                         {here + rest_cost};
          const double peak                          {std::max(size, rest_largest)};
          if (!found || total < cost || (total == cost && peak < largest)) {
            found   = true;
            cost    = total;
            largest = peak;
            path.assign(1, step);
            path.insert(path.end(), rest.begin(), rest.end());
          }
        }
      }
    }
 
    [[nodiscard]] inline std::vector<contraction_step> contraction_path(std::vector<std::string>           labels,
                                                                        const std::string&                 output,
                                                                        const std::map<char, std::size_t>& sizes)
    {
      std::vector<contraction_step> path;
      if (labels.size() <= exhaustive_path_operands) {
        double cost    {0.0};
        double largest {0.0};
        cheapest_path(labels, output, sizes, path, cost, largest);
        return path;
      }
      const auto shares {[&labels](std::size_t i, std::size_t j) {
                           return labels[i].find_first_of(labels[j]) != std::string::npos;
                         }};
      while (labels.size() > 1) {
        contraction_step best      {.first = 0, .second = 1};
        double           best_cost {-1.0};
        double           best_size {0.0};
        bool             connected {false}; // a pair sharing a label exists: an outer product waits
        for (std::size_t i {0}; i < labels.size(); ++i) {
          for (std::size_t j {i + 1}; j < labels.size(); ++j) {
            connected = connected || shares(i, j);
          }
        }
        for (std::size_t i {0}; i < labels.size(); ++i) {
          for (std::size_t j {i + 1}; j < labels.size(); ++j) {
            if (connected && !shares(i, j)) {
              continue;
            }
            const double here {extent_product(labels[i] + labels[j], sizes)};
            const double size {extent_product(pair_keeps(labels, i, j, output), sizes)};
            if (best_cost < 0.0 || here < best_cost || (here == best_cost && size < best_size)) {
              best      = {.first = i, .second = j};
              best_cost = here;
              best_size = size;
            }
          }
        }
        path.push_back(best);
        labels = after_step(labels, best, output);
      }
      return path;
    }
 
    [[nodiscard]] inline std::size_t ellipsis_in(const std::string& subscripts)
    {
      const std::size_t at   {subscripts.find("...")};
      const auto        dots {static_cast<std::size_t>(std::count(subscripts.begin(), subscripts.end(), '.'))};
      if (dots != (at == std::string::npos ? 0U : 3U)) {
        throw einsum_error("einsum: '.' stands only as one ellipsis, '...', per subscript");
      }
      return at;
    }
 
    template <typename T>
    [[nodiscard]] tensor<T> einsum_impl(const std::string&                   spec_in,
                                        const std::vector<const tensor<T>*>& operands);
 
    template <typename T>
    [[nodiscard]] tensor<T> einsum_ellipsis(const std::string&                   spec,
                                            const std::vector<const tensor<T>*>& operands)
    {
      const std::size_t        arrow  {spec.find("->")};
      std::vector<std::string> inputs {""};
      for (const char ch : spec.substr(0, arrow)) {
        if (ch == ',') {
          inputs.emplace_back();
        }
        else {
          inputs.back().push_back(ch);
        }
      }
      if (inputs.size() != operands.size()) {
        throw einsum_error("einsum: number of operands does not match the spec");
      }
      std::vector<std::size_t> at(inputs.size());
      std::vector<std::size_t> under(inputs.size(), 0); // the axes each operand's ellipsis stands for
      std::size_t              widest {0};
      for (std::size_t i {0}; i < inputs.size(); ++i) {
        at[i] = ellipsis_in(inputs[i]);
        if (at[i] != std::string::npos) {
          const std::size_t letters {inputs[i].size() - 3};
          if (operands[i]->rank() < letters) {
            throw einsum_error("einsum: operand rank does not match its subscripts");
          }
          under[i] = operands[i]->rank() - letters;
          widest   = std::max(widest, under[i]);
        }
      }
      static constexpr std::string_view alphabet {"ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz"};
      std::string                       fresh;
      for (const char letter : alphabet) {
        if (fresh.size() < widest && spec.find(letter) == std::string::npos) {
          fresh.push_back(letter);
        }
      }
      if (fresh.size() < widest) {
        throw einsum_error("einsum: too many axes under the ellipsis for the letters left");
      }
      std::vector<std::size_t> extent(widest, 1);
      for (std::size_t i {0}; i < inputs.size(); ++i) {
        for (std::size_t d {0}; d < under[i]; ++d) {
          const std::size_t dim       {operands[i]->shape()[at[i] + d]};
          std::size_t&      broadcast {extent[widest - under[i] + d]};
          if (dim != 1 && broadcast != 1 && dim != broadcast) {
            throw dimension_error("einsum: the axes under the ellipsis do not broadcast");
          }
          broadcast = dim == 1 ? broadcast : dim;
        }
      }
      std::deque<tensor<T>>         owned; // the operands with axes dropped (a deque keeps their addresses)
      std::vector<const tensor<T>*> next;
      std::string                   rewritten;
      for (std::size_t i {0}; i < inputs.size(); ++i) {
        rewritten += i == 0 ? "" : ",";
        if (at[i] == std::string::npos) {
          rewritten += inputs[i];
          next.push_back(operands[i]);
          continue;
        }
        std::string              labels {inputs[i].substr(0, at[i])};
        std::vector<std::size_t> shape(operands[i]->shape().begin(), operands[i]->shape().begin() + static_cast<std::ptrdiff_t>(at[i]));
        for (std::size_t d {0}; d < under[i]; ++d) {
          const std::size_t axis {widest - under[i] + d};
          const std::size_t dim  {operands[i]->shape()[at[i] + d]};
          if (dim != extent[axis]) { // 1 where another operand's is longer: repeated, its letter dropped
            continue;
          }
          labels.push_back(fresh[axis]);
          shape.push_back(dim);
        }
        labels += inputs[i].substr(at[i] + 3);
        shape.insert(shape.end(), operands[i]->shape().begin() + static_cast<std::ptrdiff_t>(at[i] + under[i]), operands[i]->shape().end());
        rewritten += labels;
        if (shape.size() == operands[i]->rank()) {
          next.push_back(operands[i]);
        }
        else {
          owned.emplace_back(std::move(shape), buffer<T>(operands[i]->data().begin(), operands[i]->data().end()));
          next.push_back(&owned.back());
        }
      }
      std::string output;
      if (arrow != std::string::npos) {
        output                  = spec.substr(arrow + 2);
        const std::size_t where {ellipsis_in(output)};
        if (where == std::string::npos && widest > 0) {
          throw einsum_error("einsum: the output keeps the axes under the ellipsis: '...' in it, as NumPy requires");
        }
        if (where != std::string::npos) {
          output.replace(where, 3, fresh);
        }
      }
      else { // implicit: the broadcast axes, then the letters that appear once, in order
        std::map<char, std::size_t> seen;
        for (const char label : spec) {
          if (is_label(label)) {
            ++seen[label];
          }
        }
        output = fresh;
        for (const auto& [label, count] : seen) {
          if (count == 1) {
            output.push_back(label);
          }
        }
      }
      return einsum_impl<T>(rewritten + "->" + output, next);
    }
 
    template <typename T>
    [[nodiscard]] tensor<T> einsum_impl(const std::string&                    spec_in,
                                        const std::vector<const tensor<T>*>&  operands)
    {
      // Tolerate whitespace anywhere in the spec, the way NumPy does ("ij, jk -> ik").
      std::string spec {spec_in};
      spec.erase(std::remove_if(spec.begin(), spec.end(),
                                [](unsigned char ch) { return std::isspace(ch) != 0; }),
                 spec.end());
      if (spec.find('.') != std::string::npos) {
        return einsum_ellipsis<T>(spec, operands);
      }
      if (spec.find("->") == std::string::npos) {
        // Implicit mode, as NumPy reads it: the output is every label that appears once, in ASCII order (`"ij,jk"`
        // is `"ij,jk->ik"`, `"ii"` the trace, `"ba"` the transpose).
        std::map<char, std::size_t> seen;
        for (const char label : spec) {
          if (label != ',') {
            ++seen[label];
          }
        }
        spec += "->";
        for (const auto& [label, count] : seen) {
          if (count == 1) {
            spec += label;
          }
        }
      }
      const std::size_t              arrow  {spec.find("->")};
      const std::vector<std::string> inputs {split_inputs(spec.substr(0, arrow))};
      const std::string              rhs    {spec.substr(arrow + 2)};
 
      if (operands.size() != inputs.size()) {
        throw einsum_error("einsum: number of operands does not match the spec");
      }
 
      // Each label's extent, checked for consistency across every axis it labels.
      std::map<char, std::size_t> sizes;
      for (std::size_t i {0}; i < operands.size(); ++i) {
        if (inputs[i].size() != operands[i]->rank()) {
          throw einsum_error("einsum: operand rank does not match its subscripts");
        }
        for (std::size_t axis {0}; axis < inputs[i].size(); ++axis) {
          const char        label {inputs[i][axis]};
          const std::size_t dim   {operands[i]->shape()[axis]};
          const auto        found {sizes.find(label)};
          if (found == sizes.end()) {
            sizes.emplace(label, dim);
          }
          else if (found->second != dim) {
            throw dimension_error("einsum: inconsistent extent for a repeated label");
          }
        }
      }
 
      // Output (free) labels: each must be present in the inputs and appear once.
      std::string output;
      for (const char ch : rhs) {
        if (!is_label(ch)) {
          throw einsum_error("einsum: invalid character in output subscripts");
        }
        if (sizes.find(ch) == sizes.end()) {
          throw einsum_error("einsum: output label not present in any input");
        }
        if (output.find(ch) != std::string::npos) {
          throw einsum_error("einsum: repeated output label");
        }
        output.push_back(ch);
      }
 
      // Summed labels: those that do not survive into the output (sorted, by map).
      std::string summed;
      for (const auto& entry : sizes) {
        if (output.find(entry.first) == std::string::npos) {
          summed.push_back(entry.first);
        }
      }
 
      // Three operands or more: first sum each one over the labels no other operand and not the output
      // carries, then contract them two at a time in the cheapest order (\ref contraction_path), each
      // intermediate keeping only the labels still to come. One loop over every label of `ij,jk,kl->il` runs n⁴
      // terms; pairs run 2·n³, and `i,j,k,l,m->` over 65 536 ones each runs five sums of 65 536 instead of 2⁸⁰.
      if (operands.size() > 2) {
        std::deque<tensor<T>>         owned; // the reductions and intermediates (a deque keeps their addresses)
        std::vector<const tensor<T>*> current;
        std::vector<std::string>      kept;
        for (std::size_t i {0}; i < operands.size(); ++i) {
          std::string elsewhere {output};
          for (std::size_t j {0}; j < operands.size(); ++j) {
            elsewhere += j == i ? std::string {} : inputs[j];
          }
          std::string keep;
          for (const char label : inputs[i]) {
            if (elsewhere.find(label) != std::string::npos && keep.find(label) == std::string::npos) {
              keep.push_back(label);
            }
          }
          if (keep == inputs[i]) { // nothing of its own to sum: the operand itself
            current.push_back(operands[i]);
          }
          else {
            const std::vector<const tensor<T>*> alone {operands[i]};
            owned.push_back(einsum_impl<T>(inputs[i] + "->" + keep, alone));
            current.push_back(&owned.back());
          }
          kept.push_back(std::move(keep));
        }
        for (const contraction_step step : contraction_path(kept, output, sizes)) {
          const std::string                   keep {kept.size() == 2 ? output : pair_keeps(kept, step.first, step.second, output)};
          const std::vector<const tensor<T>*> pair {current[step.first], current[step.second]};
          owned.push_back(einsum_impl<T>(kept[step.first] + "," + kept[step.second] + "->" + keep, pair));
          current.erase(current.begin() + static_cast<std::ptrdiff_t>(step.second));
          current.erase(current.begin() + static_cast<std::ptrdiff_t>(step.first));
          kept.erase(kept.begin() + static_cast<std::ptrdiff_t>(step.second));
          kept.erase(kept.begin() + static_cast<std::ptrdiff_t>(step.first));
          current.push_back(&owned.back());
          kept.push_back(keep);
        }
        return std::move(owned.back()); // the last contraction's, the output
      }
 
      // Fast path: one operand, nothing summed, each label once: a permutation of its axes (a transpose), a copy.
      if (operands.size() == 1 && summed.empty() && inputs[0].size() == output.size()) {
        std::vector<std::size_t> out_shape;
        for (const char label : output) {
          out_shape.push_back(sizes.at(label));
        }
        return tensor<T> {std::move(out_shape), inputs[0] == output ? operands[0]->data() : permuted(operands[0]->data(), inputs[0], output, sizes)};
      }
 
      // Fast path: one operand, each label once, its summed labels one run of its axes and the output the others in
      // order (ij->i, ij->, ij->j, ijk->ik): sum_over_run.
      if (operands.size() == 1 && sizes.size() == inputs[0].size()) {
        if (std::optional<tensor<T>> sums {sum_over_run(inputs[0], output, operands[0]->shape(), operands[0]->data().data())}) {
          return *std::move(sums);
        }
      }
 
      // Fast path: a plain two-operand single contraction in canonical layout —
      // `A` ending with the contracted label, `B` starting with it, and the
      // output being A's leading free labels then B's trailing ones — is exactly
      // a matrix product, routed to the packed kernel (gemm.hpp). (Any stray copy of
      // the contracted label leaves it in a "free" part, so the output would
      // carry a summed label and this equality fails — correctly declining it.)
      if (operands.size() == 2 && summed.size() == 1) {
        const char         contracted {summed.front()};
        const std::string& s0         {inputs[0]};
        const std::string& s1         {inputs[1]};
        if (!s0.empty() && !s1.empty() && s0.back() == contracted && s1.front() == contracted
            && output == s0.substr(0, s0.size() - 1) + s1.substr(1)) {
          const std::size_t        inner {sizes.at(contracted)};
          std::vector<std::size_t> out_shape;
          for (std::size_t k {0}; k + 1 < s0.size(); ++k) {
            out_shape.push_back(sizes.at(s0[k]));
          }
          for (std::size_t k {1}; k < s1.size(); ++k) {
            out_shape.push_back(sizes.at(s1[k]));
          }
          // The result's elements, 0 where one of its extents is 0 whatever the others (an empty contraction's
          // operands may hold any number along their other axes); else checked, and the rows and columns within.
          if (element_count(out_shape.begin(), out_shape.end()) == 0) {
            return tensor<T> {std::move(out_shape), buffer<T> {}};
          }
          const auto        split {out_shape.begin() + static_cast<std::ptrdiff_t>(s0.size() - 1)};
          const std::size_t rows  {element_count(out_shape.begin(), split)};
          const std::size_t cols  {element_count(split, out_shape.end())};
          return canonical_product<T>(operands[0]->data().data(), operands[1]->data().data(), rows, inner, cols, std::move(out_shape));
        }
      }
      if (operands.size() == 2) {
        if (std::optional<tensor<T>> paired {contract_pair<T>(inputs[0], *operands[0], inputs[1], *operands[1], output, summed, sizes)}) {
          return *std::move(paired);
        }
      }
 
      // Every distinct label, output (free) first then summed: this fixes the
      // index space and the row-major order of the result.
      const std::string labels   {output + summed};
      const std::size_t n_labels {labels.size()};
      const std::size_t n_free   {output.size()};
      const std::size_t n_ops    {operands.size()};
 
      std::vector<std::size_t> extent(n_labels);
      for (std::size_t k {0}; k < n_labels; ++k) {
        extent[k] = sizes.at(labels[k]);
      }
 
      // coef[i][k]: operand i's flat-offset change per unit of label k — the sum
      // of its row-major strides over the axes carrying that label (so a label
      // repeated within one operand, e.g. "ii", walks that operand's diagonal).
      std::vector<std::vector<std::size_t>> coef(n_ops, std::vector<std::size_t>(n_labels, 0));
      for (std::size_t i {0}; i < n_ops; ++i) {
        const std::vector<std::size_t>& shape  {operands[i]->shape()};
        std::size_t                     stride {1};
        for (std::size_t axis {inputs[i].size()}; axis-- > 0;) {
          coef[i][labels.find(inputs[i][axis])]  += stride;
          stride                                 *= shape[axis];
        }
      }
 
      // Fast path: peel the innermost label of each group — the last free label
      // and the last summed label — into a tight unit-of-stride inner loop, and
      // run the odometer only over the *remaining* (outer) labels. The innermost
      // axis is then a contiguous sweep amortising one odometer step over its
      // whole extent: a strided write for a transpose/broadcast, a strided
      // multiply-accumulate for a contraction — both auto-vectoriser-friendly.
      const bool        has_free_inner {n_free > 0};
      const bool        has_sum_inner  {n_labels > n_free};
      const std::size_t inner_free     {has_free_inner ? extent[n_free - 1] : 1};
      const std::size_t inner_sum      {has_sum_inner ? extent[n_labels - 1] : 1};
      const std::size_t free_outer_end {has_free_inner ? n_free - 1 : 0};
      const std::size_t sum_outer_end  {has_sum_inner ? n_labels - 1 : n_free};
 
      // Per-operand stride of each peeled axis (0 when that axis is absent).
      std::vector<std::size_t> free_step(n_ops, 0);
      std::vector<std::size_t> sum_step(n_ops, 0);
      for (std::size_t i {0}; i < n_ops; ++i) {
        if (has_free_inner) {
          free_step[i] = coef[i][n_free - 1];
        }
        if (has_sum_inner) {
          sum_step[i] = coef[i][n_labels - 1];
        }
      }
 
      std::size_t out_size {1};
      for (std::size_t k {0}; k < n_free; ++k) {
        out_size = detail::checked_mul(out_size, extent[k]); // guard the result allocation size
      }
      std::vector<std::size_t> out_shape(extent.begin(),
                                         extent.begin() + static_cast<std::ptrdiff_t>(n_free));
      if (out_size == 0) { // an empty result computes nothing, whatever its other extents multiply to
        return tensor<T> {std::move(out_shape), buffer<T> {}};
      }
      std::size_t out_outer {1}; // bounded by out_size
      for (std::size_t k {0}; k < free_outer_end; ++k) {
        out_outer *= extent[k];
      }
      std::size_t sum_outer {1};
      for (std::size_t k {n_free}; k < sum_outer_end; ++k) {
        sum_outer = detail::checked_mul(sum_outer, extent[k]); // a wrapped count would sum nothing, silently
      }
 
      std::vector<std::size_t> index(n_labels, 0); // value of each outer label
      std::vector<std::size_t> offset(n_ops, 0);   // each operand's flat offset (outer labels only)
      buffer<T>                result;
      result.reserve(out_size);
 
      for (std::size_t of {0}; of < out_outer; ++of) {
        for (std::size_t fe {0}; fe < inner_free; ++fe) {
          T acc {};
          for (std::size_t os {0}; os < sum_outer; ++os) {
            for (std::size_t se {0}; se < inner_sum; ++se) {
              // The first factor itself, then the others: a complex 1 times an element was no identity (0 · inf
              // is a NaN, and −0 − (−0) a +0), so a trace or a sum of complex infinities or signed zeros went wrong.
              T prod {operands[0]->data()[offset[0] + (fe * free_step[0]) + (se * sum_step[0])]};
              for (std::size_t i {1}; i < n_ops; ++i) {
                prod *= operands[i]->data()[offset[i] + (fe * free_step[i]) + (se * sum_step[i])];
              }
              acc += prod;
            }
            // Step the outer-summed odometer (labels [n_free, sum_outer_end)).
            for (std::size_t k {sum_outer_end}; k-- > n_free;) {
              if (++index[k] < extent[k]) {
                for (std::size_t i {0}; i < n_ops; ++i) {
                  offset[i] += coef[i][k];
                }
                break;
              }
              index[k] = 0;
              for (std::size_t i {0}; i < n_ops; ++i) {
                offset[i] -= coef[i][k] * (extent[k] - 1);
              }
            }
          }
          result.push_back(acc);
        }
        // Step the outer-free odometer (labels [0, free_outer_end)).
        for (std::size_t k {free_outer_end}; k-- > 0;) {
          if (++index[k] < extent[k]) {
            for (std::size_t i {0}; i < n_ops; ++i) {
              offset[i] += coef[i][k];
            }
            break;
          }
          index[k] = 0;
          for (std::size_t i {0}; i < n_ops; ++i) {
            offset[i] -= coef[i][k] * (extent[k] - 1);
          }
        }
      }
      return tensor<T> {std::move(out_shape), std::move(result)};
    }
  } // namespace detail
 
  template <typename T, typename ... Rest>
  [[nodiscard]] tensor<T> einsum(const std::string& spec,
                                 const tensor<T>&   first,
                                 const Rest& ...    rest)
  {
    const std::vector<const tensor<T>*> operands {&first, (&rest)...};
    return detail::einsum_impl<T>(spec, operands);
  }
 
  template <typename T>
  [[nodiscard]] tensor<T> einsum(const std::string&                   spec,
                                 const std::vector<const tensor<T>*>& operands)
  {
    return detail::einsum_impl<T>(spec, operands);
  }
 
  template <typename T>
  [[nodiscard]] tensor<T> matmul(const tensor<T>& a,
                                 const tensor<T>& b)
  {
    const std::vector<std::size_t>& sa {a.shape()};
    const std::vector<std::size_t>& sb {b.shape()};
    if (sa.empty() || sb.empty()) {
      throw dimension_error("matmul: a rank-0 operand has no matrix to multiply");
    }
    const bool        row    {sa.size() == 1};
    const bool        column {sb.size() == 1};
    const std::size_t m      {row ? 1 : sa[sa.size() - 2]};
    const std::size_t k      {sa.back()};
    const std::size_t n      {column ? 1 : sb.back()};
    if ((column ? sb[0] : sb[sb.size() - 2]) != k) {
      throw dimension_error("matmul: the inner extents of the operands differ");
    }
    const std::size_t        lead_a {sa.size() - (row ? 1 : 2)};
    const std::size_t        lead_b {sb.size() - (column ? 1 : 2)};
    const std::size_t        rank   {std::max(lead_a, lead_b)};
    std::vector<std::size_t> batch(rank);
    std::vector<std::size_t> stride_a(rank, 0); // 0 along an axis the operand repeats
    std::vector<std::size_t> stride_b(rank, 0);
    std::size_t              step_a {m * k};
    std::size_t              step_b {k * n};
    for (std::size_t d {rank}; d-- > 0;) {
      const std::size_t da {d + lead_a >= rank ? sa[d + lead_a - rank] : 1};
      const std::size_t db {d + lead_b >= rank ? sb[d + lead_b - rank] : 1};
      if (da != db && da != 1 && db != 1) {
        throw dimension_error("matmul: the leading axes of the operands do not broadcast");
      }
      batch[d]    = da == 1 ? db : da;
      stride_a[d] = da == 1 ? 0 : step_a;
      stride_b[d] = db == 1 ? 0 : step_b;
      step_a      = detail::checked_mul(step_a, da);
      step_b      = detail::checked_mul(step_b, db);
    }
    std::size_t count {1};
    for (const std::size_t extent : batch) {
      count = detail::checked_mul(count, extent);
    }
    std::vector<std::size_t> shape {batch};
    if (!row) {
      shape.push_back(m);
    }
    if (!column) {
      shape.push_back(n);
    }
    const std::array<std::size_t, 3> elements {count, m, n};
    buffer<T>                        c(detail::element_count(elements.begin(), elements.end()));
    detail::parallel_zero(c.data(), c.size()); // gemm adds to C
    if (rank == 0) {                           // no stack: one product
      detail::gemm(m, n, k, a.data().data(), k, b.data().data(), n, c.data(), n);
      return tensor<T> {std::move(shape), std::move(c)};
    }
    const auto product {[&](std::size_t p) {
                          std::size_t rest {p};
                          std::size_t offset_a {0};
                          std::size_t offset_b {0};
                          for (std::size_t d {batch.size()}; d-- > 0;) {
                            const std::size_t index {rest % batch[d]};
                            rest     /= batch[d];
                            offset_a += index * stride_a[d];
                            offset_b += index * stride_b[d];
                          }
                          detail::gemm(m, n, k, a.data().data() + offset_a, k, b.data().data() + offset_b, n, c.data() + (p * m * n), n);
                        }};
    detail::parallel_for(count, m * n * k / detail::gemm_kernel_speedup, product); // the stack's products in parallel
    return tensor<T> {std::move(shape), std::move(c)};
  }
} // namespace numerical
 
#endif // NUMERICAL_TENSOR_HPP