Source
stdlib/ndarray/ndarray.hpp
1
// Copyright (c) 2026 BigBrain LLC. MIT-licensed (see LICENSE).2
// Original work; see ACKNOWLEDGMENTS.md for the open-source ideas we build upon.3
#pragma once5
/**6
* @file ndarray.hpp7
* @brief cheatah `ndarray` — our own numpy-flavored N-dimensional array8
* (`basic_ndarray<T>` over any @ref Element type; `NDArray` is the `double`9
* default) with NumPy broadcasting, surfaced as a `NDArray` class plus free10
* functions (a .purr program writes `ndarray.zeros([2, 3])`).11
* See https://numpy.org/doc/stable/user/basics.broadcasting.html.12
*13
* `import ndarray` includes this header and links `libcheatah_ndarray`. Unit tests:14
* `stdlib/tests/ndarray_test.cpp`; the suite runs under AddressSanitizer (the `asan`15
* preset) and Valgrind (`security/run-valgrind.sh`) on every QA-gate run.16
*17
* @note Design (the "pointers + a bit of thinking"): the elements live in a shared18
* buffer (`std::shared_ptr<buffer_t<T>>`) and an array is a VIEW into19
* it — {shape, strides, offset}. That makes reshape and **broadcast**20
* zero-copy: to stretch a dimension of size 1 we give it a stride of 0, so21
* every index along it reads the same element. Shared ownership = memory-safe,22
* no manual frees. `size` below is the element count (product of dims).23
*/24
#include <algorithm>25
#include <array>26
#include <cmath>27
#include <complex>28
#include <concepts>29
#include <cstddef>30
#include <initializer_list>31
#include <limits>32
#include <memory>33
#include <new> // placement new, for the default-init buffer allocator34
#include <numeric>35
#include <sstream>36
#include <stdexcept>37
#include <string>38
#include <type_traits>39
#include <utility> // std::forward / std::move40
#include <version> // __cpp_lib_execution feature-test macro41
#include <vector>43
// The unsequenced execution policy lets std::transform/std::reduce vectorize. libstdc++44
// provides it; Apple's libc++ historically ships no usable <execution>, so guard on the45
// feature-test macro and fall back to the plain (policy-less) overloads where it's46
// absent. This is speed-neutral: `unseq` is unsequenced (no threads, no TBB) and for47
// these simple element-wise loops the -O3 -march=native auto-vectorizer produces the48
// same SIMD either way — the transcendental vectorization comes from ufunc_simd.cpp's49
// libmvec/Accelerate kernels, not from this policy.50
#if defined(__cpp_lib_execution)51
#include <execution>52
#define CHEATAH_UNSEQ std::execution::unseq,53
#else54
#define CHEATAH_UNSEQ55
#endif57
namespace cheatah::ndarray {59
/// Numeric<T>: an arithmetic element type an ndarray can store (int or float60
/// family). Storage, construction, and elementwise +-* require only this.61
template <typename T>62
concept Numeric = std::is_arithmetic_v<T>;63
/// FloatingPoint<T>: a real floating type. The linalg decompositions (solve, inv,64
/// det, svd, eig) need division/√, so they constrain to this — calling them on an65
/// integer array fails with a clear "FloatingPoint not satisfied", not template spam.66
template <typename T>67
concept FloatingPoint = std::floating_point<T>;69
/// @cond INTERNAL70
template <typename T>71
struct is_complex : std::false_type {};72
template <typename U>73
struct is_complex<std::complex<U>> : std::bool_constant<std::is_floating_point_v<U>> {};74
/// @endcond76
/// Whether `T` is a `std::complex` of a floating type — the trait behind @ref Field.77
template <typename T>78
inline constexpr bool is_complex_v = is_complex<T>::value;80
/// Field<T>: a scalar an ndarray can store — a real arithmetic type OR a81
/// `std::complex` of a floating type. This is what makes **complex** matrices and82
/// vectors first-class (Hermitian operators, complex wavefunctions), and lets a83
/// REAL matrix yield the COMPLEX eigenvalues it mathematically has.84
template <typename T>85
concept Field = std::is_arithmetic_v<T> || is_complex_v<T>;87
/// Element<T>: the broadest bound — any type an ndarray may STORE. A @ref Field (real/complex88
/// number) OR any MOVABLE type, so a fixed-size struct (a 2-D point, an RGBA colour, a GPU89
/// vertex) lives in an ndarray too. Elements are MOVED into the buffer on construction and the90
/// backing buffer is never silently deep-copied (copying an ndarray shares the buffer — an O(1)91
/// view). The arithmetic surface (elementwise ops, ufuncs, reductions, linalg) stays constrained92
/// to @ref Field and the duplicating factories to @ref Copyable, so a move-only element still93
/// stores / indexes / views / moves — it simply cannot be summed or deep-copied, and the compiler94
/// says so by design (cheatah discourages hidden copies on hot data).95
template <typename T>96
concept Element = Field<T> || std::movable<T>;98
/// Copyable<T>: an @ref Element that may ALSO be duplicated. It gates only the value-fill99
/// factories (full / full_like) and reshape's deep copy — the paths that replicate elements.100
/// Numbers and ordinary copyable POD structs satisfy it; a move-only struct does NOT, so an101
/// accidental deep copy of it fails to compile. Copying an ndarray CONTAINER never needs this —102
/// it is always a shared-buffer view (and the GPU borrows that buffer in place, never copying it).103
template <typename T>104
concept Copyable = Element<T> && std::copyable<T>;106
/// Subscript<T>: what may address an axis — an integer, OR a scoped `enum class` whose ordinal names107
/// the position (a column label). This concept is the ONLY door through which a scoped enum becomes an108
/// integer: `enum class` values stay strongly typed everywhere else, and the implicit109
/// enum-to-index conversion is confined to array subscripting, exactly where a named column belongs.110
template <typename T>111
concept Subscript = std::is_convertible_v<T, long long> || std::is_enum_v<T>;113
/// The integer position an @ref Subscript addresses. For a scoped enum this is its underlying ordinal;114
/// this `static_cast` is the whole of the enum-to-index conversion the language sanctions.115
/// @tparam Ix the subscript type: an integer, or a scoped `enum class`.116
/// @param i the subscript to resolve.117
/// @return the integer position it names (a scoped enum's underlying ordinal).118
/// @complexity O(1). @alloc none.119
/// @test Fixarray.EnumIndexingOnVectorsAndMatrices120
template <Subscript Ix>121
[[nodiscard]] constexpr long long subscript_index(Ix i) noexcept {122
return static_cast<long long>(i);123
}125
/// @cond INTERNAL126
template <typename T>127
struct real_base {128
using type = T;129
};130
template <typename U>131
struct real_base<std::complex<U>> {132
using type = U;133
};134
/// @endcond136
/// The real type underlying a @ref Field `T` (`double` for both `double` and137
/// `complex<double>`).138
template <typename T>139
using real_base_t = typename real_base<T>::type;141
/// complex_of_t<T>: the complex type over T's real base. `eig`/`eigvals` return an142
/// array of these, because a real matrix can have complex eigenvalues (conjugate143
/// pairs) — e.g. the rotation matrix [[0,-1],[1,0]] has eigenvalues ±i.144
template <typename T>145
using complex_of_t = std::complex<real_base_t<T>>;147
namespace detail {148
/// @cond INTERNAL149
/// An allocator identical to `std::allocator<T>` in every respect EXCEPT that150
/// DEFAULT (no-value) construction — what `vector(n)` / `resize(n)` perform — leaves a151
/// trivially-constructible element UNINITIALIZED instead of value-initializing it to 0.152
///153
/// Every ndarray op that allocates a result buffer it then overwrites in full (binary154
/// ops, ufuncs, reshape, array()) would otherwise pay a throwaway zero-fill of the whole155
/// buffer first. That wasted write pass is hidden on compute-heavy ops but DOMINATES156
/// bandwidth-bound ones — `add` was ≈1.5× of NumPy purely from the extra memset. With157
/// this allocator the sizing path skips it; the value-filling forms (`assign(n, v)`,158
/// `vector(n, v)` used by zeros/full/scalar) are untouched and still initialize.159
template <typename T>160
struct default_init_allocator : std::allocator<T> {161
using std::allocator<T>::allocator;162
template <typename U>163
struct rebind {164
using other = default_init_allocator<U>;165
};166
/// Default construction: a trivially-constructible element is left uninitialized.167
template <typename U>168
void construct(U* p) noexcept(std::is_nothrow_default_constructible_v<U>) {169
::new (static_cast<void*>(p)) U; // default-init (no `()`): no zeroing for trivial U170
}171
/// Every other construction (value-fill, copy, emplace) behaves exactly as normal.172
template <typename U, typename... Args>173
void construct(U* p, Args&&... args) {174
::new (static_cast<void*>(p)) U(std::forward<Args>(args)...);175
}176
};177
/// @endcond179
/// C-order (row-major) strides for a shape.180
inline std::vector<std::ptrdiff_t> contiguous_strides(const std::vector<std::size_t>& shape) {181
std::vector<std::ptrdiff_t> s(shape.size());182
std::ptrdiff_t step = 1;183
for (std::size_t i = shape.size(); i-- > 0;) {184
s[i] = step;185
step *= static_cast<std::ptrdiff_t>(shape[i]);186
}187
return s;188
}189
/// Overflow-checked product of the dimensions (a wrapped size_t would under-allocate,190
/// turning later element access into out-of-bounds writes — reject it up front).191
inline std::size_t product(const std::vector<std::size_t>& shape) {192
std::size_t t = 1;193
for (std::size_t d : shape) {194
if (d != 0 && t > std::numeric_limits<std::size_t>::max() / d) {195
throw std::runtime_error("ndarray: shape too large (size overflow)");196
}197
t *= d;198
}199
return t;200
}201
/// Convert signed dims/indices to sizes, rejecting negatives (a negative cast to202
/// size_t becomes huge -> under-allocation / OOB). Validate at the boundary.203
inline std::vector<std::size_t> to_size(const std::vector<long long>& v) {204
std::vector<std::size_t> out(v.size());205
for (std::size_t i = 0; i < v.size(); ++i) {206
if (v[i] < 0) throw std::runtime_error("ndarray: negative dimension or index");207
out[i] = static_cast<std::size_t>(v[i]);208
}209
return out;210
}211
/// Advance a C-order multi-index odometer; false when it wraps past the end.212
inline bool next_index(std::vector<std::size_t>& idx, const std::vector<std::size_t>& shape) {213
for (std::size_t i = shape.size(); i-- > 0;) {214
if (++idx[i] < shape[i]) return true;215
idx[i] = 0;216
}217
return false;218
}220
/// Peel `std::vector<>` layers off a (possibly deeply nested) list type to reach the221
/// leaf scalar — `nested_scalar_t<std::vector<std::vector<double>>>` is `double`.222
template <typename V> struct nested_scalar { using type = V; };223
template <typename U> struct nested_scalar<std::vector<U>> {224
using type = typename nested_scalar<U>::type;225
};226
template <typename V> using nested_scalar_t = typename nested_scalar<V>::type;228
/// Whether `V` is a `std::vector<…>` (used to tell a nested list from a scalar leaf).229
template <typename V> inline constexpr bool is_std_vector_v = false;230
template <typename U> inline constexpr bool is_std_vector_v<std::vector<U>> = true;232
/// Flatten a scalar leaf into the C-order buffer (recursion base case).233
template <Element T>234
void nested_collect(T x, std::vector<T>& flat, std::vector<std::size_t>&, std::size_t) {235
flat.push_back(std::move(x));236
}237
/// Walk a nested list: record each axis length the first time it is seen, reject a238
/// ragged list (a row whose length differs from its siblings — numpy does too), and239
/// flatten the leaves in C-order. The leaf scalar must be a @ref Field.240
template <typename U>241
requires Field<nested_scalar_t<U>>242
void nested_collect(const std::vector<U>& v, std::vector<nested_scalar_t<U>>& flat,243
std::vector<std::size_t>& shape, std::size_t depth) {244
if (depth == shape.size()) shape.push_back(v.size());245
else if (shape[depth] != v.size())246
throw std::runtime_error("ndarray: array(...) ragged nested list (a row's length "247
"differs from its siblings)");248
for (const U& e : v) nested_collect(e, flat, shape, depth + 1);249
}250
} // namespace detail252
/// The backing store of an ndarray: a flat, contiguous, shared element buffer. It uses253
/// @ref detail::default_init_allocator so a freshly-sized result buffer that an op is254
/// about to overwrite in full is not needlessly zero-filled first. zeros/full/scalar,255
/// which value-fill, are unaffected — only the no-value sizing path skips initialization.256
template <Element T>257
using buffer_t = std::vector<T, detail::default_init_allocator<T>>;259
/**260
* @brief An N-dimensional array of `T` (a @ref Field element type — real or complex):261
* a view ({shape, strides, offset}) over a shared element buffer.262
*263
* Copies are cheap and share the buffer; reshape/broadcast produce new views without264
* copying elements. Index math goes through @ref at, which bounds-checks. The element265
* type is deduced from the data (e.g. `array([1,2,3])` is integer, `array([1.0,…])`266
* is double); `NDArray` is the default `basic_ndarray<double>`. Complex element types267
* (`std::complex<double>`) make complex matrices/vectors — and the complex eigenvalues268
* a real matrix can have — first-class.269
*/270
template <Element T>271
class basic_ndarray {272
public:273
using value_type = T; ///< The stored element type (an @ref Element `T`).274
/**275
* Construct an empty 0-d array with a fresh empty buffer.276
*277
* Leaves shape and strides empty and the buffer holding no elements; note this278
* is distinct from a 0-d scalar (see @ref scalar), whose buffer holds one element.279
* @complexity O(1).280
* @alloc allocates the empty shared buffer.281
* @test CheatahNDArray.ToStringScalar282
* @systest StdlibE2E.Ndarray283
*/284
basic_ndarray() : data_(std::make_shared<buffer_t<T>>()) {}285
/**286
* Construct a contiguous array of @p shape filled with @p fill.287
*288
* Allocates a fresh buffer of `product(shape)` elements all set to @p fill and289
* computes C-order (row-major) strides; the dimension product is overflow-checked.290
* @param shape the dimensions.291
* @param fill value for every element.292
* @complexity O(size).293
* @alloc allocates a new shared buffer (`shared_ptr<buffer_t<T>>`) of294
* `product(shape)` elements; throws if the shape overflows size_t.295
* @test CheatahNDArray.ShapeFactoriesAndReductions296
* @systest StdlibE2E.Ndarray297
*/298
explicit basic_ndarray(std::vector<std::size_t> shape, T fill = T{}) // contiguous299
: data_(std::make_shared<buffer_t<T>>()), shape_(std::move(shape)) {300
// resize (default-init: no zero pass) then std::fill — the fill goes through301
// operator= on built elements, which keeps libstdc++'s memset/SIMD fast path.302
// (vector::assign(n, v) would route through the allocator's construct, which a303
// default-init allocator forces element-by-element — measurably slower.)304
data_->resize(detail::product(shape_));305
std::fill(data_->begin(), data_->end(), fill);306
strides_ = detail::contiguous_strides(shape_);307
}308
/**309
* Build a contiguous array of @p shape whose buffer is sized but left310
* UNINITIALIZED — for internal ops (binary ops, ufuncs, reshape, array) that311
* immediately write every element, so paying for a zero-fill first is pure waste.312
* @param shape the dimensions.313
* @return an array of @p shape with an uninitialized buffer.314
* @complexity O(1) beyond the allocation (no element initialization).315
* @alloc allocates an uninitialized `product(shape)`-element buffer; overflow-checked.316
* @test CheatahNDArray.BroadcastingAdd317
* @systest StdlibE2E.Ndarray318
*/319
static basic_ndarray uninitialized(std::vector<std::size_t> shape) {320
basic_ndarray out;321
out.shape_ = std::move(shape);322
out.data_->resize(detail::product(out.shape_)); // default-init alloc -> no zero-fill323
out.strides_ = detail::contiguous_strides(out.shape_);324
return out;325
}326
/**327
* Construct a view from explicit buffer/shape/strides/offset (used by views).328
*329
* Stores the supplied members verbatim with no validation or copy, so the new330
* array shares ownership of @p data; callers (e.g. @ref broadcast_to) are331
* responsible for passing strides/offset that stay within the buffer.332
* @param data the shared element buffer.333
* @param shape the dimensions.334
* @param strides element strides per dimension.335
* @param offset starting flat offset into @p data.336
* @complexity O(1).337
* @alloc shares @p data, no element copy.338
* @test CheatahNDArray.BroadcastTo339
* @systest StdlibE2E.Ndarray340
*/341
basic_ndarray(std::shared_ptr<buffer_t<T>> data, std::vector<std::size_t> shape,342
std::vector<std::ptrdiff_t> strides, std::size_t offset)343
: data_(std::move(data)), shape_(std::move(shape)), strides_(std::move(strides)),344
offset_(offset) {}346
/**347
* The shape (dimensions).348
* @return reference to the shape vector.349
* @complexity O(1).350
* @alloc none.351
* @test CheatahNDArray.ShapeFactoriesAndReductions352
* @systest StdlibE2E.Ndarray353
*/354
const std::vector<std::size_t>& shape() const { return shape_; }355
/**356
* The element strides.357
* @return reference to the strides vector.358
* @complexity O(1).359
* @alloc none.360
* @test CheatahNDArray.BroadcastTo361
* @systest StdlibE2E.Ndarray362
*/363
const std::vector<std::ptrdiff_t>& strides() const { return strides_; }364
/**365
* The number of dimensions (rank).366
* @return `shape().size()`.367
* @complexity O(1).368
* @alloc none.369
* @test CheatahNDArray.BroadcastingAdd370
* @systest StdlibE2E.Ndarray371
*/372
std::size_t ndim() const { return shape_.size(); }373
/**374
* The element count (product of dims; 1 for a 0-d scalar).375
*376
* Recomputes the overflow-checked product of the shape on each call rather than377
* caching it; an empty (no-dimension) shape yields 1.378
* @return the number of elements.379
* @complexity O(ndim).380
* @alloc none; throws on size overflow.381
* @test CheatahNDArray.ShapeFactoriesAndReductions382
* @systest StdlibE2E.Ndarray383
*/384
std::size_t size() const { return detail::product(shape_); } // 1 for a 0-d scalar386
/**387
* MUTABLE element reference by multi-index — the write path behind cheatah388
* subscript assignment `x[i] = v` / `x[i, j] = v`. Negative indices count389
* from the end of their dimension; rank and bounds are checked.390
* @param ixs one (possibly negative) coordinate per dimension.391
* @return a writable reference into the backing buffer.392
* @complexity O(ndim). @alloc none.393
* @test CheatahNDArray.SubscriptReadWrite394
* @crtest LangFeatures.NdarraySubscript395
* @systest StdlibE2E.Ndarray396
*/397
template <typename... Ix>398
requires((Subscript<Ix> && ...) && sizeof...(Ix) > 0)399
T& item_ref(Ix... ixs) {400
const std::array<long long, sizeof...(Ix)> raw{subscript_index(ixs)...};401
if (raw.size() != shape_.size())402
throw std::out_of_range("ndarray subscript: wrong number of indices");403
std::size_t pos = offset_;404
for (std::size_t d = 0; d < raw.size(); ++d) {405
long long i = raw[d];406
const long long n = static_cast<long long>(shape_[d]);407
if (i < 0) i += n;408
if (i < 0 || i >= n) throw std::out_of_range("ndarray subscript out of range");409
pos += static_cast<std::size_t>(i) * strides_[d];410
}411
return (*data_)[pos];412
}414
/// 1-D subscript: `mask[i] = v` assignment and raw element writes. Accepts an integer or a scoped415
/// enum column label (see @ref Subscript), so `q[QuatCol::W] = v` writes column W of a 1-D row.416
/// @param i the element index (negative counts from the end, Python-style).417
/// @return a mutable reference to element @p i.418
/// @complexity O(1). @alloc none.419
/// @test CheatahNDArray.SubscriptReadWrite420
template <Subscript Ix>421
T& operator[](Ix i) {422
return item_ref(subscript_index(i));423
}425
/**426
* Read one element by a full multidimensional index (one component per axis), resolved via427
* the array's strides. Computes the flat buffer position as428
* `offset + sum(index[i] * strides[i])`, so it correctly resolves views (including429
* broadcast dims with stride 0). Bounds- and rank-checked.430
* @param index the per-axis indices; its size must equal the array's rank.431
* @return a copy of the addressed element.432
* @throws std::runtime_error if @p index has the wrong rank or any component is out of range.433
* @complexity O(rank).434
* @alloc none.435
* @test CheatahNDArray.AtVectorRankAndRangeErrors,436
* CheatahNDArray.RejectsMaliciousShapesAndIndices437
* @systest StdlibE2E.Ndarray438
*/439
T at(const std::vector<std::size_t>& index) const { // element via strides440
// Bounds-check: a wrong-rank or out-of-range index would otherwise compute an441
// offset outside the backing buffer (out-of-bounds read).442
if (index.size() != shape_.size()) {443
throw std::runtime_error("ndarray: index has the wrong number of dimensions");444
}445
std::ptrdiff_t off = static_cast<std::ptrdiff_t>(offset_);446
for (std::size_t i = 0; i < index.size(); ++i) {447
if (index[i] >= shape_[i]) throw std::runtime_error("ndarray: index out of range");448
off += static_cast<std::ptrdiff_t>(index[i]) * strides_[i];449
}450
return (*data_)[static_cast<std::size_t>(off)];451
}452
/**453
* The shared backing buffer.454
* @return reference to the element buffer shared_ptr.455
* @complexity O(1).456
* @alloc none.457
* @test CheatahNDArray.BroadcastingAdd458
* @systest StdlibE2E.Ndarray459
*/460
const std::shared_ptr<buffer_t<T>>& buffer() const { return data_; }461
/**462
* The flat offset into the buffer where this view starts.463
* @return the offset.464
* @complexity O(1).465
* @alloc none.466
* @test CheatahNDArray.BroadcastTo467
* @systest StdlibE2E.Ndarray468
*/469
std::size_t offset() const { return offset_; }471
/**472
* Python-style text rendering, e.g. `"[[1, 2], [3, 4]]"` — the `str()` hook that473
* makes an NDArray printable via `io.print` (io's `HasStr` protocol). Defers to474
* the free `to_string`; defined out-of-line below, where `to_string` is declared.475
* @return the array formatted as nested brackets.476
* @complexity O(n) in the element count.477
* @alloc allocates the result string.478
* @test CheatahNDArray.PrettyPrintAbbreviatesLarge479
* @systest StdlibE2E.Ndarray480
*/481
std::string str() const;483
/**484
* Pretty-print hook used by `io.print`: renders the array in nested-bracket form but485
* ABBREVIATES a large array with `...` (numpy-style edge items), so printing a big array486
* stays readable. `io.rprint` (and `str()`/`operator<<`) keep the FULL untruncated form —487
* slice the array and `rprint` the subset to see everything.488
* @param os destination stream.489
* @param indent unused (arrays are self-delimiting with brackets); present so `io.print`490
* detects this hook uniformly with struct pretty-printers.491
* @complexity O(printed elements).492
* @alloc allocates intermediate strings.493
* @test CheatahNDArray.PrettyPrintAbbreviatesLarge494
* @systest StdlibE2E.Ndarray495
*/496
void cheatah_pretty_print(std::ostream& os, long long indent) const;498
private:499
std::shared_ptr<buffer_t<T>> data_;500
std::vector<std::size_t> shape_;501
std::vector<std::ptrdiff_t> strides_; // element strides502
std::size_t offset_ = 0;503
};505
/// The default ndarray element type is `double` — `NDArray` names that506
/// instantiation (the std::string ↔ std::basic_string<char> pattern), so existing507
/// code and the linalg routines keep working unchanged.508
using NDArray = basic_ndarray<double>;510
/**511
* The broadcast result shape of two shapes (NumPy rules).512
*513
* Aligns the shapes from the trailing (rightmost) dimension, treating missing514
* leading dims as 1; each output dim is the non-1 input dim, and two unequal dims515
* that are both not 1 are incompatible and throw.516
* @param a first shape.517
* @param b second shape.518
* @return the broadcast shape (trailing-aligned).519
* @complexity O(max(ndim)).520
* @alloc allocates the small result vector; throws if the shapes are incompatible.521
* @test CheatahNDArray.BroadcastShapeRules522
* @systest StdlibE2E.Ndarray523
*/524
std::vector<std::size_t> broadcast_shapes(const std::vector<std::size_t>& a,525
const std::vector<std::size_t>& b);527
// ==========================================================================528
// Templated free functions. The element type T is deduced from the data529
// (array([1,2,3]) -> long long, array([1.0,…]) -> double); every op is530
// constrained by Numeric. Elementwise ops vectorize via the std::execution531
// policies (declarative SIMD); broadcasting/strided views fall back to a532
// C-order scalar walk. NDArray (= basic_ndarray<double>) is the default.533
// ==========================================================================535
/**536
* Whether @p a is a contiguous C-order block (no broadcast/stride-0/permuted view),537
* so its elements live consecutively from `offset()` and can be walked flatly.538
* @param a the array (or view) to test.539
* @return true if @p a's strides are the C-order strides for its shape.540
* @complexity O(ndim).541
* @alloc none.542
* @test CheatahNDArray.BroadcastingAdd543
* @systest StdlibE2E.Ndarray544
*/545
template <Element T>546
inline bool is_contiguous(const basic_ndarray<T>& a) {547
// C-order contiguity WITHOUT materializing the reference strides: walk the dims548
// back-to-front and check each stride equals the running size product. The old549
// `strides() == contiguous_strides(shape())` heap-allocated a vector on every call550
// — a fixed cost that dominated small-n reductions (e.g. dot, where it ran twice551
// per call). Same result as the comparison, zero allocation. O(ndim).552
const std::vector<std::size_t>& shape = a.shape();553
const std::vector<std::ptrdiff_t>& strides = a.strides();554
std::ptrdiff_t expect = 1;555
for (std::size_t i = shape.size(); i-- > 0;) {556
if (strides[i] != expect) return false;557
expect *= static_cast<std::ptrdiff_t>(shape[i]);558
}559
return true;560
}562
/**563
* A zero-copy view of @p a stretched to @p target (size-1 / missing dims get stride 0).564
* @param a source array.565
* @param target the shape to stretch to.566
* @return a VIEW sharing @p a's buffer (no element copy).567
* @complexity O(rank of @p target).568
* @alloc no element copy — only the view's stride vector; throws if the shapes are569
* not broadcast-compatible.570
* @test CheatahNDArray.BroadcastTo571
* @systest StdlibE2E.Ndarray572
*/573
template <Element T>574
basic_ndarray<T> broadcast_to(const basic_ndarray<T>& a, const std::vector<std::size_t>& target) {575
const std::size_t n = target.size();576
if (a.ndim() > n) throw std::runtime_error("ndarray: cannot broadcast to fewer dimensions");577
std::vector<std::ptrdiff_t> ns(n, 0); // stretched / missing dims -> stride 0578
const std::size_t pad = n - a.ndim();579
for (std::size_t i = 0; i < a.ndim(); ++i) {580
const std::size_t adim = a.shape()[i];581
if (adim == target[pad + i]) {582
ns[pad + i] = a.strides()[i];583
} else if (adim != 1) {584
throw std::runtime_error("ndarray: shape not broadcastable to target");585
} // adim == 1 -> stride stays 0 (stretch)586
}587
return basic_ndarray<T>(a.buffer(), target, ns, a.offset());588
}590
// ---- factories (shapes arrive from cheatah as list[int]) ----591
/**592
* 1-D array from a list of values; the element type is the list's element type593
* (`array([1,2,3])` is integer, `array([1.0,…])` is double).594
* @param values the elements, copied into a fresh contiguous buffer.595
* @return a contiguous 1-D `basic_ndarray<T>`.596
* @complexity O(n).597
* @alloc allocates a new buffer of `values.size()` elements.598
* @test CheatahNDArray.ShapeFactoriesAndReductions599
* @crtest NdarrayCompileRun.Array600
* @systest StdlibE2E.Ndarray601
*/602
template <Copyable T>603
requires (!detail::is_std_vector_v<T>) // a vector-of-vectors is a NESTED list (overload below)604
basic_ndarray<T> array(const std::vector<T>& values) {605
basic_ndarray<T> a = basic_ndarray<T>::uninitialized({values.size()});606
std::copy(values.begin(), values.end(), a.buffer()->begin());607
return a;608
}609
/**610
* 1-D array that MOVES its elements out of @p values into a fresh buffer (no element copy) — the611
* no-copy build path, and the ONLY `array` overload a move-only element type has. A temporary612
* `array(std::vector<T>{…})` binds here automatically; a named lvalue you want to keep uses the613
* copying overload above.614
* @param values the elements, moved into a fresh contiguous buffer (left moved-from).615
* @return a contiguous 1-D `basic_ndarray<T>`.616
* @complexity O(n).617
* @alloc allocates a new buffer of `values.size()` elements.618
* @test CheatahNDArray.ArrayMoveIn619
* @systest StdlibE2E.Ndarray620
*/621
template <Element T>622
requires (!detail::is_std_vector_v<T>) // a vector-of-vectors is a NESTED list (overload below)623
basic_ndarray<T> array(std::vector<T>&& values) {624
basic_ndarray<T> a = basic_ndarray<T>::uninitialized({values.size()});625
std::move(values.begin(), values.end(), a.buffer()->begin());626
return a;627
}628
/**629
* N-dimensional array from a **nested** list — `array([[1, 2], [3, 4]])` is 2-D,630
* `array([[[1],[2]],[[3],[4]]])` is 3-D, and so on to any depth. The shape is read off631
* the nesting (outer list = axis 0, …) and the leaf scalar type is deduced; the list632
* must be **rectangular** (every sibling row the same length) or it throws, exactly as633
* numpy rejects a ragged array. Selected only when the argument is itself a list of634
* lists, so it never competes with the 1-D @ref array overload above.635
* @tparam V the element type of the outer list — itself a `std::vector<…>` whose leaf636
* is a @ref Field.637
* @param values the nested list (rows, planes, …), copied into a fresh C-order buffer.638
* @return a contiguous `basic_ndarray<T>` of the inferred shape.639
* @complexity O(size).640
* @alloc allocates one buffer of `product(shape)` elements; throws on a ragged list.641
* @test CheatahNDArray.NestedArrayConstruction642
* @crtest NdarrayCompileRun.NestedArray643
* @systest StdlibE2E.Ndarray644
*/645
template <typename V>646
requires detail::is_std_vector_v<V> && Element<detail::nested_scalar_t<V>>647
basic_ndarray<detail::nested_scalar_t<V>> array(const std::vector<V>& values) {648
using T = detail::nested_scalar_t<V>;649
std::vector<std::size_t> shape;650
std::vector<T> flat;651
detail::nested_collect(values, flat, shape, 0);652
basic_ndarray<T> a = basic_ndarray<T>::uninitialized(shape);653
std::move(flat.begin(), flat.end(), a.buffer()->begin());654
return a;655
}656
/**657
* `array({1, 2, 3})` — braced-list overload (deduces T from the initializer_list,658
* which the `std::vector<T>` overload can't do directly).659
* @param values the elements as a braced list.660
* @return a contiguous 1-D `basic_ndarray<T>`.661
* @complexity O(n).662
* @alloc allocates a new buffer.663
* @test CheatahNDArray.ShapeFactoriesAndReductions664
* @systest StdlibE2E.Ndarray665
*/666
template <Copyable T>667
requires (!detail::is_std_vector_v<T>) // nested braces route to the nested-list overload668
basic_ndarray<T> array(std::initializer_list<T> values) {669
return array(std::vector<T>(values));670
}671
/**672
* 0-D scalar array (broadcasts to anything); element type deduced from @p value.673
* @param value the single element.674
* @return a 0-d `basic_ndarray<T>`.675
* @complexity O(1).676
* @alloc allocates a one-element buffer.677
* @test CheatahNDArray.ElementwiseAndScalarBroadcast678
* @crtest NdarrayCompileRun.Scalar679
* @systest StdlibE2E.Ndarray680
*/681
template <Copyable T>682
basic_ndarray<T> scalar(T value) {683
basic_ndarray<T> a; // 0-d684
a.buffer()->assign(1, value);685
return a;686
}687
/**688
* Array of @p shape filled with 0 (a `double` array by default; rejects negatives).689
* @param shape the dimensions (signed; throws on a negative).690
* @return a zero-filled `NDArray`.691
* @complexity O(size).692
* @alloc allocates a new buffer; throws on negative/overflowing dims.693
* @test CheatahNDArray.ShapeFactoriesAndReductions694
* @crtest NdarrayCompileRun.Zeros695
* @systest StdlibE2E.Ndarray696
*/697
inline NDArray zeros(const std::vector<long long>& shape) {698
return NDArray(detail::to_size(shape), 0.0);699
}700
/**701
* Array of @p shape filled with 1 (a `double` array by default; rejects negatives).702
* @param shape the dimensions (signed; throws on a negative).703
* @return a one-filled `NDArray`.704
* @complexity O(size).705
* @alloc allocates a new buffer; throws on negative/overflowing dims.706
* @test CheatahNDArray.ShapeFactoriesAndReductions707
* @crtest NdarrayCompileRun.Ones708
* @systest StdlibE2E.Ndarray709
*/710
inline NDArray ones(const std::vector<long long>& shape) {711
return NDArray(detail::to_size(shape), 1.0);712
}713
/**714
* Array of @p shape filled with @p value; element type deduced from @p value.715
* @param shape the dimensions (signed; throws on a negative).716
* @param value the fill value (its type is the array's element type).717
* @return a filled `basic_ndarray<T>`.718
* @complexity O(size).719
* @alloc allocates a new buffer; throws on negative/overflowing dims.720
* @test CheatahNDArray.RejectsMaliciousShapesAndIndices721
* @crtest NdarrayCompileRun.Full722
* @systest StdlibE2E.Ndarray723
*/724
template <Copyable T>725
basic_ndarray<T> full(const std::vector<long long>& shape, T value) {726
return basic_ndarray<T>(detail::to_size(shape), value);727
}728
/**729
* A fresh array with the SAME shape and element type as @p a, filled with @p value730
* (≈ `numpy.full_like`). The companion `zeros_like` / `ones_like` default the fill.731
* @param a the array whose shape and element type to mirror.732
* @param value the fill value.733
* @return a same-shape `basic_ndarray<T>` filled with @p value.734
* @complexity O(size).735
* @alloc allocates a new buffer.736
* @test CheatahNDArray.LikeFactories737
*/738
template <Copyable T>739
basic_ndarray<T> full_like(const basic_ndarray<T>& a, T value) {740
return basic_ndarray<T>(a.shape(), value);741
}742
/**743
* A zero-filled array with the SAME shape and element type as @p a (≈ `numpy.zeros_like`) —744
* the idiomatic way to allocate a matching gradient/velocity/scratch buffer for an existing array.745
* @param a the array whose shape and element type to mirror.746
* @return a same-shape `basic_ndarray<T>` of zeros.747
* @complexity O(size).748
* @alloc allocates a new buffer.749
* @test CheatahNDArray.LikeFactories750
*/751
template <Copyable T>752
basic_ndarray<T> zeros_like(const basic_ndarray<T>& a) {753
return basic_ndarray<T>(a.shape(), T{});754
}755
/**756
* A one-filled array with the SAME shape and element type as @p a (≈ `numpy.ones_like`).757
* @param a the array whose shape and element type to mirror.758
* @return a same-shape `basic_ndarray<T>` of ones.759
* @complexity O(size).760
* @alloc allocates a new buffer.761
* @test CheatahNDArray.LikeFactories762
*/763
template <Copyable T>764
basic_ndarray<T> ones_like(const basic_ndarray<T>& a) {765
return basic_ndarray<T>(a.shape(), T{1});766
}767
/**768
* 1-D range `[start, stop)` stepping by @p step; element type deduced from the args.769
* @param start first value.770
* @param stop exclusive bound.771
* @param step increment (throws if zero); a step pointing away from @p stop yields empty.772
* @return a 1-D `basic_ndarray<T>` of the generated values.773
* @complexity O(count).774
* @alloc allocates a new buffer (built via a growing temporary vector, then copied);775
* throws if @p step is zero.776
* @test CheatahNDArray.Arange777
* @crtest NdarrayCompileRun.Arange778
* @systest StdlibE2E.Ndarray779
*/780
template <Numeric T>781
basic_ndarray<T> arange(T start, T stop, T step) {782
if (step == T{}) throw std::runtime_error("ndarray: arange step must be non-zero");783
std::vector<T> v;784
for (T x = start; (step > T{}) ? (x < stop) : (x > stop); x += step) v.push_back(x);785
return array(v);786
}787
/**788
* Reshape @p a to @p shape (same element count); reads in C-order so views/broadcasts789
* are flattened into a fresh contiguous buffer (a copy, not an alias).790
* @param a source array.791
* @param shape the new dimensions (signed; throws on a negative).792
* @return a new contiguous `basic_ndarray<T>` with the data in C-order.793
* @complexity O(size).794
* @alloc allocates a new buffer; throws on size mismatch or negative dims.795
* @test CheatahNDArray.BroadcastingAdd, CheatahNDArray.ReshapeSizeMismatchThrows796
* @crtest NdarrayCompileRun.Reshape797
* @systest StdlibE2E.Ndarray798
*/799
template <Copyable T>800
basic_ndarray<T> reshape(const basic_ndarray<T>& a, const std::vector<long long>& shape) {801
const std::vector<std::size_t> ns = detail::to_size(shape);802
if (detail::product(ns) != a.size()) {803
throw std::runtime_error("ndarray: cannot reshape, size mismatch");804
}805
basic_ndarray<T> out = basic_ndarray<T>::uninitialized(ns); // every element is written below806
auto& buf = *out.buffer();807
// Contiguous source (the common case — e.g. reshaping a freshly built array): copy808
// the flat block in one shot instead of walking a per-element bounds-checked809
// odometer.810
if (is_contiguous(a)) {811
const T* src = a.buffer()->data() + a.offset();812
std::copy(src, src + a.size(), buf.begin());813
return out;814
}815
std::vector<std::size_t> idx(a.ndim(), 0);816
std::size_t flat = 0;817
do {818
buf[flat++] = a.at(idx);819
} while (a.ndim() != 0 && detail::next_index(idx, a.shape()));820
return out;821
}822
/**823
* Convert @p a to a new array with element type @p U — numpy's `a.astype(dtype)`. Every element824
* is `static_cast` into @p U, so this is the way to build a NARROW-element array (a smaller memory825
* footprint): `array([1,2,3]).astype(i16)` is a `basic_ndarray<std::int16_t>` — 2 bytes/element,826
* not 8. Reads @p a in C-order (a view/broadcast is flattened into a fresh contiguous buffer — a827
* copy, never an alias), same shape out as in. Widening is exact; narrowing truncates/wraps at the828
* target width (as in C / a numpy fixed dtype). Constrained to conversions that actually exist829
* (`convertible_to`), so e.g. complex→real fails with a clear concept error, not template spam.830
* @tparam U the destination element type (the only type spelled at the call site).831
* @param a source array (any @ref Field element type convertible to @p U).832
* @return a fresh contiguous `basic_ndarray<U>` of @p a's shape.833
* @complexity O(size).834
* @alloc allocates the result buffer.835
* @test CheatahNDArray.AstypeNarrowsAndWidens836
* @crtest NdarrayCompileRun.Astype837
* @systest StdlibE2E.Ndarray838
*/839
template <Field U, Field T>840
requires std::convertible_to<T, U>841
basic_ndarray<U> astype(const basic_ndarray<T>& a) {842
basic_ndarray<U> out = basic_ndarray<U>::uninitialized(a.shape()); // every element is written below843
auto& buf = *out.buffer();844
if (is_contiguous(a)) { // contiguous source: one straight cast pass, no odometer845
const T* src = a.buffer()->data() + a.offset();846
for (std::size_t i = 0; i < a.size(); ++i) buf[i] = static_cast<U>(src[i]);847
return out;848
}849
std::vector<std::size_t> idx(a.ndim(), 0);850
std::size_t flat = 0;851
do {852
buf[flat++] = static_cast<U>(a.at(idx));853
} while (a.ndim() != 0 && detail::next_index(idx, a.shape()));854
return out;855
}857
// ---- element-wise ops (broadcasting, vectorized) ----858
/// @cond INTERNAL — implementation plumbing / compiler-selected reuse overloads (README documents the public forms)859
/**860
* Broadcast @p a and @p b to their common shape and apply @p op elementwise into a861
* fresh contiguous result. Fast path: when both operands are contiguous, a flat862
* `std::transform` under the `unseq` policy (SIMD); otherwise a C-order scalar walk.863
* @param a first operand.864
* @param b second operand.865
* @param op the binary operation applied to corresponding elements.866
* @return `op(a, b)` broadcast to the common shape; throws if shapes don't broadcast.867
* @complexity O(size of result).868
* @alloc allocates the result buffer.869
* @test CheatahNDArray.BroadcastingAdd870
* @systest StdlibE2E.Ndarray871
*/872
template <Field T, typename Op>873
basic_ndarray<T> binary_op(const basic_ndarray<T>& a, const basic_ndarray<T>& b, Op op) {874
const std::vector<std::size_t> rshape = broadcast_shapes(a.shape(), b.shape());875
basic_ndarray<T> out = basic_ndarray<T>::uninitialized(rshape); // every element written below876
auto& obuf = *out.buffer();877
// Scalar fast paths: `array ⊕ scalar` (or the reverse) is by far the most common878
// broadcast, and the general strided walk below does a bounds-checked at() per879
// element (no SIMD). When the other operand is a single value over a contiguous880
// full-shape array, it's a flat loop we hand to the unseq transform so it vectorizes881
// the same way the array⊕array path does.882
if (b.size() == 1 && a.shape() == rshape && is_contiguous(a)) {883
const T s = (*b.buffer())[b.offset()];884
const auto first = a.buffer()->begin() + a.offset();885
std::transform(CHEATAH_UNSEQ first, first + obuf.size(), obuf.begin(),886
[s, op](T x) { return op(x, s); });887
return out;888
}889
if (a.size() == 1 && b.shape() == rshape && is_contiguous(b)) {890
const T s = (*a.buffer())[a.offset()];891
const auto first = b.buffer()->begin() + b.offset();892
std::transform(CHEATAH_UNSEQ first, first + obuf.size(), obuf.begin(),893
[s, op](T x) { return op(s, x); });894
return out;895
}896
const basic_ndarray<T> av = broadcast_to(a, rshape);897
const basic_ndarray<T> bv = broadcast_to(b, rshape);898
if (is_contiguous(av) && is_contiguous(bv)) {899
const auto& abuf = *av.buffer();900
const auto& bbuf = *bv.buffer();901
std::transform(CHEATAH_UNSEQ abuf.begin() + av.offset(),902
abuf.begin() + av.offset() + obuf.size(), bbuf.begin() + bv.offset(),903
obuf.begin(), op);904
return out;905
}906
std::vector<std::size_t> idx(rshape.size(), 0);907
std::size_t flat = 0;908
do {909
obuf[flat++] = op(av.at(idx), bv.at(idx));910
} while (!rshape.empty() && detail::next_index(idx, rshape));911
return out;912
}914
/**915
* Elementwise `out = op(a, b)` (broadcasting) into the CALLER'S buffer @p out — the user-provided-output916
* form of binary_op, NO allocation: a hot loop hands the same scratch array every call. @p out must917
* already hold the broadcast result shape and be contiguous; it MAY alias a full-shape operand (the write918
* is index-local, so `add(x, x, y)` is fine).919
* @param out the destination (mutated; must be contiguous and match the broadcast shape).920
* @param a first operand.921
* @param b second operand.922
* @param op the binary combiner.923
* @complexity O(size of out).924
* @alloc none.925
* @test CheatahNDArray.BinaryOpIntoReusesBuffer926
*/927
template <Field T, typename Op>928
void binary_op_into(basic_ndarray<T>& out, const basic_ndarray<T>& a, const basic_ndarray<T>& b, Op op) {929
const std::vector<std::size_t> rshape = broadcast_shapes(a.shape(), b.shape());930
if (out.shape() != rshape || !is_contiguous(out)) {931
throw std::invalid_argument(932
"ndarray binary op (out form): out must be contiguous and match the broadcast shape");933
}934
auto& obuf = *out.buffer();935
const auto odst = obuf.begin() + static_cast<std::ptrdiff_t>(out.offset());936
if (b.size() == 1 && a.shape() == rshape && is_contiguous(a)) { // array ⊕ scalar937
const T s = (*b.buffer())[b.offset()];938
const auto af = a.buffer()->begin() + static_cast<std::ptrdiff_t>(a.offset());939
std::transform(CHEATAH_UNSEQ af, af + static_cast<std::ptrdiff_t>(out.size()), odst,940
[s, op](T x) { return op(x, s); });941
return;942
}943
if (a.size() == 1 && b.shape() == rshape && is_contiguous(b)) { // scalar ⊕ array944
const T s = (*a.buffer())[a.offset()];945
const auto bf = b.buffer()->begin() + static_cast<std::ptrdiff_t>(b.offset());946
std::transform(CHEATAH_UNSEQ bf, bf + static_cast<std::ptrdiff_t>(out.size()), odst,947
[s, op](T x) { return op(s, x); });948
return;949
}950
const basic_ndarray<T> av = broadcast_to(a, rshape);951
const basic_ndarray<T> bv = broadcast_to(b, rshape);952
if (is_contiguous(av) && is_contiguous(bv)) { // both full-shape contiguous: flat SIMD transform953
const auto af = av.buffer()->begin() + static_cast<std::ptrdiff_t>(av.offset());954
std::transform(CHEATAH_UNSEQ af, af + static_cast<std::ptrdiff_t>(out.size()),955
bv.buffer()->begin() + static_cast<std::ptrdiff_t>(bv.offset()), odst, op);956
return;957
}958
std::vector<std::size_t> idx(rshape.size(), 0); // strided operand fallback (still no alloc for out)959
std::size_t flat = 0;960
do {961
odst[static_cast<std::ptrdiff_t>(flat++)] = op(av.at(idx), bv.at(idx));962
} while (!rshape.empty() && detail::next_index(idx, rshape));963
}964
/// @endcond966
// Shared elementwise combiners: ONE functor type per op, used by BOTH the allocating967
// forms (add/sub/mul/divide) and the in-place compound operators (+=/-=/*=//=). Using a968
// single type means `binary_op` is instantiated once per op rather than once per call969
// site, so the in-place fallback reuses the same (already-tested) instantiation instead970
// of a duplicate whose scalar/contiguous fast paths are unreachable through it.971
namespace detail {972
struct add_op { template <typename T> T operator()(T x, T y) const { return x + y; } };973
struct sub_op { template <typename T> T operator()(T x, T y) const { return x - y; } };974
struct mul_op { template <typename T> T operator()(T x, T y) const { return x * y; } };975
struct div_op { template <typename T> T operator()(T x, T y) const { return x / y; } };976
// Reversed combiners: reuse the RIGHT operand in place for non-commutative ops, i.e. compute977
// `dst = src OP dst` so that `a - std::move(b)` / `a / std::move(b)` can write through b's buffer.978
struct rsub_op { template <typename T> T operator()(T x, T y) const { return y - x; } };979
struct rdiv_op { template <typename T> T operator()(T x, T y) const { return y / x; } };980
} // namespace detail982
/**983
* Element-wise `a + b` with broadcasting.984
* @param a first operand.985
* @param b second operand.986
* @return `a + b` broadcast to the common shape.987
* @complexity O(size of result). @alloc allocates the result.988
* @test CheatahNDArray.BroadcastingAdd989
* @crtest NdarrayCompileRun.Add990
* @systest StdlibE2E.Ndarray991
*/992
template <Field T>993
basic_ndarray<T> add(const basic_ndarray<T>& a, const basic_ndarray<T>& b) {994
return binary_op(a, b, detail::add_op{});995
}996
/**997
* Element-wise `a - b` with broadcasting.998
* @param a first operand.999
* @param b second operand.1000
* @return `a - b` broadcast to the common shape.1001
* @complexity O(size of result). @alloc allocates the result.1002
* @test CheatahNDArray.ElementwiseAndScalarBroadcast1003
* @crtest NdarrayCompileRun.Sub1004
* @systest StdlibE2E.Ndarray1005
*/1006
template <Field T>1007
basic_ndarray<T> sub(const basic_ndarray<T>& a, const basic_ndarray<T>& b) {1008
return binary_op(a, b, detail::sub_op{});1009
}1010
/**1011
* Element-wise `a * b` with broadcasting.1012
* @param a first operand.1013
* @param b second operand.1014
* @return `a * b` broadcast to the common shape.1015
* @complexity O(size of result). @alloc allocates the result.1016
* @test CheatahNDArray.ElementwiseAndScalarBroadcast1017
* @crtest NdarrayCompileRun.Mul1018
* @systest StdlibE2E.Ndarray1019
*/1020
template <Field T>1021
basic_ndarray<T> mul(const basic_ndarray<T>& a, const basic_ndarray<T>& b) {1022
return binary_op(a, b, detail::mul_op{});1023
}1024
/**1025
* Element-wise `a / b` with broadcasting (an integer element type does integer division).1026
* @param a numerator.1027
* @param b denominator (float division follows IEEE-754: /0 yields inf/nan, no throw).1028
* @return `a / b` broadcast to the common shape.1029
* @complexity O(size of result). @alloc allocates the result.1030
* @test CheatahNDArray.ElementwiseAndScalarBroadcast1031
* @crtest NdarrayCompileRun.Divide1032
* @systest StdlibE2E.Ndarray1033
*/1034
template <Field T>1035
basic_ndarray<T> divide(const basic_ndarray<T>& a, const basic_ndarray<T>& b) {1036
return binary_op(a, b, detail::div_op{});1037
}1039
/// @cond INTERNAL — implementation plumbing / compiler-selected reuse overloads (README documents the public forms)1040
// ---- user-provided-output forms: write into the caller's buffer, NO allocation (see binary_op_into) ----1041
/**1042
* Element-wise `a + b` into the caller's buffer @p out (out FIRST) — the buffer-reuse overload, so a1043
* hot loop hands the same scratch every call. @p out must be contiguous with the broadcast shape; it1044
* may alias a full-shape operand (the write is index-local).1045
* @param out destination, overwritten.1046
* @param a,b operands (broadcastable to @p out's shape).1047
* @complexity O(size of result). @alloc none.1048
* @test CheatahNDArray.BinaryOpIntoReusesBuffer1049
*/1050
template <Field T>1051
void add(basic_ndarray<T>& out, const basic_ndarray<T>& a, const basic_ndarray<T>& b) {1052
binary_op_into(out, a, b, detail::add_op{});1053
}1054
/**1055
* Element-wise `a - b` into @p out (out FIRST), no allocation. @see add(out, a, b).1056
* @param out destination, overwritten.1057
* @param a,b operands (broadcastable to @p out's shape).1058
*/1059
template <Field T>1060
void sub(basic_ndarray<T>& out, const basic_ndarray<T>& a, const basic_ndarray<T>& b) {1061
binary_op_into(out, a, b, detail::sub_op{});1062
}1063
/**1064
* Element-wise `a * b` into @p out (out FIRST), no allocation. @see add(out, a, b).1065
* @param out destination, overwritten.1066
* @param a,b operands (broadcastable to @p out's shape).1067
*/1068
template <Field T>1069
void mul(basic_ndarray<T>& out, const basic_ndarray<T>& a, const basic_ndarray<T>& b) {1070
binary_op_into(out, a, b, detail::mul_op{});1071
}1072
/**1073
* Element-wise `a / b` into @p out (out FIRST), no allocation. @see add(out, a, b).1074
* @param out destination, overwritten.1075
* @param a,b operands (broadcastable to @p out's shape).1076
*/1077
template <Field T>1078
void divide(basic_ndarray<T>& out, const basic_ndarray<T>& a, const basic_ndarray<T>& b) {1079
binary_op_into(out, a, b, detail::div_op{});1080
}1081
/// @endcond1083
// ---- Infix operators & in-place compound assignment ------------------------1084
// cheatah lowers `a + b` / `a * 2.0` / `a += b` on ndarrays straight to these1085
// C++ operators. Infix forms are the elementwise free functions (broadcasting1086
// included); a bare arithmetic scalar on either side is wrapped via scalar().1087
// The compound forms mutate the LEFT OPERAND'S BUFFER IN PLACE on the common1088
// (contiguous) layout — no allocation, so a hot loop can reuse one array for1089
// an entire run — falling back to the allocating elementwise path only for1090
// non-contiguous views or true broadcasts.1092
/// @cond INTERNAL — implementation plumbing / compiler-selected reuse overloads (README documents the public forms)1093
/**1094
* In-place elementwise update `a = op(a, b)`, writing through @p a's buffer.1095
* Contiguous @p a with a same-shape contiguous or single-element @p b runs as1096
* a flat vectorizable transform with NO allocation; anything else falls back1097
* to the allocating binary_op and rebinds @p a to the result.1098
* @param a the destination array (mutated).1099
* @param b the right operand (same shape, or a single element).1100
* @param op the elementwise combiner.1101
* @complexity O(size of @p a).1102
* @alloc none on the contiguous fast path; one result array on the fallback.1103
* @test CheatahNDArray.CompoundAssignInPlace1104
* @crtest LangFeatures.NdarrayOperators1105
* @systest StdlibE2E.Ndarray1106
*/1107
template <typename T, typename Op>1108
void compound_apply(basic_ndarray<T>& a, const basic_ndarray<T>& b, Op op) {1109
if (is_contiguous(a) && (b.size() == 1 || b.shape() == a.shape())) {1110
auto& abuf = *a.buffer();1111
const std::size_t n = a.size();1112
const auto first = abuf.begin() + static_cast<std::ptrdiff_t>(a.offset());1113
if (b.size() == 1) {1114
const T s = (*b.buffer())[b.offset()];1115
std::transform(CHEATAH_UNSEQ first, first + static_cast<std::ptrdiff_t>(n), first,1116
[s, op](T x) { return op(x, s); });1117
return;1118
}1119
if (is_contiguous(b)) {1120
const auto& bbuf = *b.buffer();1121
std::transform(CHEATAH_UNSEQ first, first + static_cast<std::ptrdiff_t>(n),1122
bbuf.begin() + static_cast<std::ptrdiff_t>(b.offset()), first, op);1123
return;1124
}1125
}1126
a = binary_op(a, b, op); // broadcast / non-contiguous fallback1127
}1128
/// @endcond1130
/// @cond INTERNAL — implementation plumbing / compiler-selected reuse overloads (README documents the public forms)1131
// ---- rvalue-reuse ("move") forms -----------------------------------------------------------------1132
// "Copy vs move" for ndarray math. `a + b` on NAMED (lvalue) arrays MUST allocate — it can't clobber1133
// `a`, which the caller still holds. But when the LEFT operand is an RVALUE — a temporary the caller1134
// has already given up: the `a + b` inside a chain `a + b + c`, or an explicit `std::move(a)` — these1135
// compute IN PLACE into that buffer and move it out: NO allocation. Selected by value category, so a1136
// buffer is only ever reused when it is safe to (no flag, no surprise mutation). Reuses the in-place1137
// compound_apply, so the same contiguous fast path / broadcast fallback / tests apply.1138
/**1139
* Element-wise `a + b` reusing the expiring left operand @p a in place (no allocation).1140
* @param a the expiring left operand; computed into and moved out.1141
* @param b the right operand (broadcastable to @p a's shape).1142
* @return the sum, in @p a's reused buffer.1143
* @complexity O(size of @p a). @alloc none on the contiguous fast path.1144
*/1145
template <Field T>1146
basic_ndarray<T> add(basic_ndarray<T>&& a, const basic_ndarray<T>& b) {1147
compound_apply(a, b, detail::add_op{});1148
return std::move(a);1149
}1150
/**1151
* Element-wise `a - b` reusing the expiring left operand @p a in place (no allocation).1152
* @param a the expiring left operand; computed into and moved out.1153
* @param b the right operand (broadcastable to @p a's shape).1154
* @return the difference, in @p a's reused buffer.1155
* @complexity O(size of @p a). @alloc none on the contiguous fast path.1156
*/1157
template <Field T>1158
basic_ndarray<T> sub(basic_ndarray<T>&& a, const basic_ndarray<T>& b) {1159
compound_apply(a, b, detail::sub_op{});1160
return std::move(a);1161
}1162
/**1163
* Element-wise `a * b` reusing the expiring left operand @p a in place (no allocation).1164
* @param a the expiring left operand; computed into and moved out.1165
* @param b the right operand (broadcastable to @p a's shape).1166
* @return the product, in @p a's reused buffer.1167
* @complexity O(size of @p a). @alloc none on the contiguous fast path.1168
*/1169
template <Field T>1170
basic_ndarray<T> mul(basic_ndarray<T>&& a, const basic_ndarray<T>& b) {1171
compound_apply(a, b, detail::mul_op{});1172
return std::move(a);1173
}1174
/**1175
* Element-wise `a / b` reusing the expiring left operand @p a in place (no allocation).1176
* @param a the expiring left operand; computed into and moved out.1177
* @param b the right operand (broadcastable to @p a's shape).1178
* @return the quotient, in @p a's reused buffer.1179
* @complexity O(size of @p a). @alloc none on the contiguous fast path.1180
*/1181
template <Field T>1182
basic_ndarray<T> divide(basic_ndarray<T>&& a, const basic_ndarray<T>& b) {1183
compound_apply(a, b, detail::div_op{});1184
return std::move(a);1185
}1187
// Right-operand reuse: when only the RIGHT operand is the expiring temporary, compute through ITS1188
// buffer instead. `+`/`*` are commutative so `op(b, a)` is the same value; `-`/`/` use the reversed1189
// combiners (`b = a OP b`). This makes `a + std::move(b)` reuse a buffer exactly like1190
// `std::move(a) + b` — the two are symmetric, neither allocates.1191
/**1192
* Element-wise `a + b` reusing the expiring right operand @p b in place (no allocation).1193
* @param a the left operand (a const lvalue the caller keeps).1194
* @param b the expiring right operand; computed into and moved out.1195
* @return the sum, in @p b's reused buffer.1196
* @complexity O(size of @p b). @alloc none on the contiguous fast path.1197
*/1198
template <Field T>1199
basic_ndarray<T> add(const basic_ndarray<T>& a, basic_ndarray<T>&& b) {1200
compound_apply(b, a, detail::add_op{});1201
return std::move(b);1202
}1203
/**1204
* Element-wise `a - b` reusing the expiring right operand @p b in place via the reversed combiner1205
* `b = a - b` (no allocation).1206
* @param a the left operand (a const lvalue the caller keeps).1207
* @param b the expiring right operand; computed into and moved out.1208
* @return the difference, in @p b's reused buffer.1209
* @complexity O(size of @p b). @alloc none on the contiguous fast path.1210
*/1211
template <Field T>1212
basic_ndarray<T> sub(const basic_ndarray<T>& a, basic_ndarray<T>&& b) {1213
compound_apply(b, a, detail::rsub_op{});1214
return std::move(b);1215
}1216
/**1217
* Element-wise `a * b` reusing the expiring right operand @p b in place (no allocation).1218
* @param a the left operand (a const lvalue the caller keeps).1219
* @param b the expiring right operand; computed into and moved out.1220
* @return the product, in @p b's reused buffer.1221
* @complexity O(size of @p b). @alloc none on the contiguous fast path.1222
*/1223
template <Field T>1224
basic_ndarray<T> mul(const basic_ndarray<T>& a, basic_ndarray<T>&& b) {1225
compound_apply(b, a, detail::mul_op{});1226
return std::move(b);1227
}1228
/**1229
* Element-wise `a / b` reusing the expiring right operand @p b in place via the reversed combiner1230
* `b = a / b` (no allocation).1231
* @param a the numerator (a const lvalue the caller keeps).1232
* @param b the expiring denominator; computed into and moved out.1233
* @return the quotient, in @p b's reused buffer.1234
* @complexity O(size of @p b). @alloc none on the contiguous fast path.1235
*/1236
template <Field T>1237
basic_ndarray<T> divide(const basic_ndarray<T>& a, basic_ndarray<T>&& b) {1238
compound_apply(b, a, detail::rdiv_op{});1239
return std::move(b);1240
}1242
// Both operands expiring: prefer reusing the LEFT (matches the chain `a + b + c`, where the left is1243
// the running accumulator). Disambiguates the otherwise-ambiguous `std::move(a) OP std::move(b)`.1244
/**1245
* Element-wise `a + b` when both operands are expiring; reuses the left buffer @p a (no allocation).1246
* @param a the expiring left operand; reused for the result.1247
* @param b the expiring right operand.1248
* @return the sum, in @p a's reused buffer.1249
* @complexity O(size of @p a). @alloc none on the contiguous fast path.1250
*/1251
template <Field T>1252
basic_ndarray<T> add(basic_ndarray<T>&& a, basic_ndarray<T>&& b) { return add(std::move(a), b); }1253
/**1254
* Element-wise `a - b` when both operands are expiring; reuses the left buffer @p a (no allocation).1255
* @param a the expiring left operand; reused for the result.1256
* @param b the expiring right operand.1257
* @return the difference, in @p a's reused buffer.1258
* @complexity O(size of @p a). @alloc none on the contiguous fast path.1259
*/1260
template <Field T>1261
basic_ndarray<T> sub(basic_ndarray<T>&& a, basic_ndarray<T>&& b) { return sub(std::move(a), b); }1262
/**1263
* Element-wise `a * b` when both operands are expiring; reuses the left buffer @p a (no allocation).1264
* @param a the expiring left operand; reused for the result.1265
* @param b the expiring right operand.1266
* @return the product, in @p a's reused buffer.1267
* @complexity O(size of @p a). @alloc none on the contiguous fast path.1268
*/1269
template <Field T>1270
basic_ndarray<T> mul(basic_ndarray<T>&& a, basic_ndarray<T>&& b) { return mul(std::move(a), b); }1271
/**1272
* Element-wise `a / b` when both operands are expiring; reuses the left buffer @p a (no allocation).1273
* @param a the expiring numerator; reused for the result.1274
* @param b the expiring denominator.1275
* @return the quotient, in @p a's reused buffer.1276
* @complexity O(size of @p a). @alloc none on the contiguous fast path.1277
*/1278
template <Field T>1279
basic_ndarray<T> divide(basic_ndarray<T>&& a, basic_ndarray<T>&& b) { return divide(std::move(a), b); }1280
/// @endcond1282
/// Elementwise infix forms: `a + b`, `a - b`, `a * b`, `a / b` (broadcasting).1283
/**1284
* Elementwise `a + b` with broadcasting (infix form of add()).1285
* @param a first operand. @param b second operand.1286
* @return the broadcast sum (a fresh array). @complexity O(size of result). @alloc the result.1287
* @test CheatahNDArray.RvalueOperandReusesBuffer1288
*/1289
template <typename T>1290
basic_ndarray<T> operator+(const basic_ndarray<T>& a, const basic_ndarray<T>& b) { return add(a, b); }1291
/**1292
* Elementwise `a - b` with broadcasting (infix form of sub()).1293
* @param a first operand. @param b second operand.1294
* @return the broadcast difference (a fresh array). @complexity O(size of result). @alloc the result.1295
*/1296
template <typename T>1297
basic_ndarray<T> operator-(const basic_ndarray<T>& a, const basic_ndarray<T>& b) { return sub(a, b); }1298
/**1299
* Elementwise `a * b` with broadcasting (infix form of mul()).1300
* @param a first operand. @param b second operand.1301
* @return the broadcast product (a fresh array). @complexity O(size of result). @alloc the result.1302
*/1303
template <typename T>1304
basic_ndarray<T> operator*(const basic_ndarray<T>& a, const basic_ndarray<T>& b) { return mul(a, b); }1305
/**1306
* Elementwise `a / b` with broadcasting (infix form of divide()).1307
* @param a numerator. @param b denominator.1308
* @return the broadcast quotient (a fresh array). @complexity O(size of result). @alloc the result.1309
* @test CheatahNDArray.DivideInfixLvalueForm1310
*/1311
template <typename T>1312
basic_ndarray<T> operator/(const basic_ndarray<T>& a, const basic_ndarray<T>& b) { return divide(a, b); }1314
/// @cond INTERNAL — implementation plumbing / compiler-selected reuse overloads (README documents the public forms)1315
/// rvalue-reuse infix forms: whichever operand is the expiring temporary is computed into IN PLACE1316
/// (no alloc). `std::move(a) + b` reuses `a`, `a + std::move(b)` reuses `b` — symmetric. A chain1317
/// `a + b + c` allocates once (for `a + b`) instead of twice. If BOTH are temporaries the left wins.1318
/** `a + b` reusing the expiring left operand @p a in place. @param a expiring left operand. @param b right operand. @return the sum, in @p a's buffer. @complexity O(size). @alloc none on the fast path. */1319
template <typename T>1320
basic_ndarray<T> operator+(basic_ndarray<T>&& a, const basic_ndarray<T>& b) { return add(std::move(a), b); }1321
/** `a - b` reusing the expiring left operand @p a in place. @param a expiring left operand. @param b right operand. @return the difference, in @p a's buffer. @complexity O(size). @alloc none on the fast path. */1322
template <typename T>1323
basic_ndarray<T> operator-(basic_ndarray<T>&& a, const basic_ndarray<T>& b) { return sub(std::move(a), b); }1324
/** `a * b` reusing the expiring left operand @p a in place. @param a expiring left operand. @param b right operand. @return the product, in @p a's buffer. @complexity O(size). @alloc none on the fast path. */1325
template <typename T>1326
basic_ndarray<T> operator*(basic_ndarray<T>&& a, const basic_ndarray<T>& b) { return mul(std::move(a), b); }1327
/** `a / b` reusing the expiring left operand @p a in place. @param a expiring numerator. @param b denominator. @return the quotient, in @p a's buffer. @complexity O(size). @alloc none on the fast path. */1328
template <typename T>1329
basic_ndarray<T> operator/(basic_ndarray<T>&& a, const basic_ndarray<T>& b) { return divide(std::move(a), b); }1330
/** `a + b` reusing the expiring right operand @p b in place. @param a left operand. @param b expiring right operand. @return the sum, in @p b's buffer. @complexity O(size). @alloc none on the fast path. */1331
template <typename T>1332
basic_ndarray<T> operator+(const basic_ndarray<T>& a, basic_ndarray<T>&& b) { return add(a, std::move(b)); }1333
/** `a - b` reusing the expiring right operand @p b in place. @param a left operand. @param b expiring right operand. @return the difference, in @p b's buffer. @complexity O(size). @alloc none on the fast path. */1334
template <typename T>1335
basic_ndarray<T> operator-(const basic_ndarray<T>& a, basic_ndarray<T>&& b) { return sub(a, std::move(b)); }1336
/** `a * b` reusing the expiring right operand @p b in place. @param a left operand. @param b expiring right operand. @return the product, in @p b's buffer. @complexity O(size). @alloc none on the fast path. */1337
template <typename T>1338
basic_ndarray<T> operator*(const basic_ndarray<T>& a, basic_ndarray<T>&& b) { return mul(a, std::move(b)); }1339
/** `a / b` reusing the expiring right operand @p b in place. @param a numerator. @param b expiring denominator. @return the quotient, in @p b's buffer. @complexity O(size). @alloc none on the fast path. */1340
template <typename T>1341
basic_ndarray<T> operator/(const basic_ndarray<T>& a, basic_ndarray<T>&& b) { return divide(a, std::move(b)); }1342
/** `a + b` with both operands expiring; reuses the left buffer @p a. @param a expiring left operand. @param b expiring right operand. @return the sum, in @p a's buffer. @complexity O(size). @alloc none on the fast path. */1343
template <typename T>1344
basic_ndarray<T> operator+(basic_ndarray<T>&& a, basic_ndarray<T>&& b) { return add(std::move(a), std::move(b)); }1345
/** `a - b` with both operands expiring; reuses the left buffer @p a. @param a expiring left operand. @param b expiring right operand. @return the difference, in @p a's buffer. @complexity O(size). @alloc none on the fast path. */1346
template <typename T>1347
basic_ndarray<T> operator-(basic_ndarray<T>&& a, basic_ndarray<T>&& b) { return sub(std::move(a), std::move(b)); }1348
/** `a * b` with both operands expiring; reuses the left buffer @p a. @param a expiring left operand. @param b expiring right operand. @return the product, in @p a's buffer. @complexity O(size). @alloc none on the fast path. */1349
template <typename T>1350
basic_ndarray<T> operator*(basic_ndarray<T>&& a, basic_ndarray<T>&& b) { return mul(std::move(a), std::move(b)); }1351
/** `a / b` with both operands expiring; reuses the left buffer @p a. @param a expiring numerator. @param b expiring denominator. @return the quotient, in @p a's buffer. @complexity O(size). @alloc none on the fast path. */1352
template <typename T>1353
basic_ndarray<T> operator/(basic_ndarray<T>&& a, basic_ndarray<T>&& b) { return divide(std::move(a), std::move(b)); }1354
/// @endcond1356
/// Scalar infix forms, either side: `a * 2.0`, `0.5 * a`, `a + 1`, `1.0 / a`, ...1357
/// The scalar converts to the array's element type (int literals work on float arrays). These go1358
/// straight to the ALLOCATING binary_op (not the reuse-enabled add/sub/...): the `scalar(s)` temporary1359
/// is 0-d, so letting it bind a buffer-reuse overload would compute the result into the scalar and1360
/// collapse it to 0-d. The array operand here is a const lvalue (the caller keeps it), so the result1361
/// must be a fresh array of the BROADCAST shape regardless.1362
/** `a + s`: add scalar @p s to every element. @param a the array. @param s the arithmetic scalar (converted to @p T). @return a fresh array of @p a's shape. @complexity O(size of @p a). @alloc the result. */1363
template <typename T, typename S>1364
requires std::is_arithmetic_v<S>1365
basic_ndarray<T> operator+(const basic_ndarray<T>& a, S s) { return binary_op(a, scalar(static_cast<T>(s)), detail::add_op{}); }1366
/** `s + a`: add scalar @p s to every element. @param s the arithmetic scalar (converted to @p T). @param a the array. @return a fresh array of @p a's shape. @complexity O(size of @p a). @alloc the result. */1367
template <typename T, typename S>1368
requires std::is_arithmetic_v<S>1369
basic_ndarray<T> operator+(S s, const basic_ndarray<T>& a) { return binary_op(scalar(static_cast<T>(s)), a, detail::add_op{}); }1370
/** `a - s`: subtract scalar @p s from every element. @param a the array. @param s the arithmetic scalar (converted to @p T). @return a fresh array of @p a's shape. @complexity O(size of @p a). @alloc the result. */1371
template <typename T, typename S>1372
requires std::is_arithmetic_v<S>1373
basic_ndarray<T> operator-(const basic_ndarray<T>& a, S s) { return binary_op(a, scalar(static_cast<T>(s)), detail::sub_op{}); }1374
/** `s - a`: elementwise @p s minus each element. @param s the arithmetic scalar (converted to @p T). @param a the array. @return a fresh array of @p a's shape. @complexity O(size of @p a). @alloc the result. @test CheatahNDArray.ScalarTimesSizeOneArrayKeepsShape */1375
template <typename T, typename S>1376
requires std::is_arithmetic_v<S>1377
basic_ndarray<T> operator-(S s, const basic_ndarray<T>& a) { return binary_op(scalar(static_cast<T>(s)), a, detail::sub_op{}); }1378
/** `a * s`: multiply every element by scalar @p s. @param a the array. @param s the arithmetic scalar (converted to @p T). @return a fresh array of @p a's shape. @complexity O(size of @p a). @alloc the result. @test CheatahNDArray.CompoundAssignInPlace */1379
template <typename T, typename S>1380
requires std::is_arithmetic_v<S>1381
basic_ndarray<T> operator*(const basic_ndarray<T>& a, S s) { return binary_op(a, scalar(static_cast<T>(s)), detail::mul_op{}); }1382
/** `s * a`: multiply every element by scalar @p s. @param s the arithmetic scalar (converted to @p T). @param a the array. @return a fresh array of @p a's shape. @complexity O(size of @p a). @alloc the result. @test CheatahNDArray.ScalarTimesSizeOneArrayKeepsShape */1383
template <typename T, typename S>1384
requires std::is_arithmetic_v<S>1385
basic_ndarray<T> operator*(S s, const basic_ndarray<T>& a) { return binary_op(scalar(static_cast<T>(s)), a, detail::mul_op{}); }1386
/** `a / s`: divide every element by scalar @p s. @param a the array. @param s the arithmetic scalar divisor (converted to @p T). @return a fresh array of @p a's shape. @complexity O(size of @p a). @alloc the result. */1387
template <typename T, typename S>1388
requires std::is_arithmetic_v<S>1389
basic_ndarray<T> operator/(const basic_ndarray<T>& a, S s) { return binary_op(a, scalar(static_cast<T>(s)), detail::div_op{}); }1390
/** `s / a`: elementwise @p s divided by each element. @param s the arithmetic scalar numerator (converted to @p T). @param a the array of divisors. @return a fresh array of @p a's shape. @complexity O(size of @p a). @alloc the result. */1391
template <typename T, typename S>1392
requires std::is_arithmetic_v<S>1393
basic_ndarray<T> operator/(S s, const basic_ndarray<T>& a) { return binary_op(scalar(static_cast<T>(s)), a, detail::div_op{}); }1395
/// @cond INTERNAL — implementation plumbing / compiler-selected reuse overloads (README documents the public forms)1396
/// rvalue-reuse scalar forms: a temporary array operand is computed into IN PLACE (no alloc). Only the1397
/// commutative `s + a` / `s * a` get a scalar-LEFT reuse form; `s - a` / `s / a` keep the allocating1398
/// const& form (a reversed in-place would be needed, not worth it for that rare case).1399
/** `a + s` reusing the expiring array @p a in place. @param a expiring array operand. @param s the arithmetic scalar (converted to @p T). @return the result, in @p a's buffer. @complexity O(size of @p a). @alloc none on the fast path. */1400
template <typename T, typename S>1401
requires std::is_arithmetic_v<S>1402
basic_ndarray<T> operator+(basic_ndarray<T>&& a, S s) { return add(std::move(a), scalar(static_cast<T>(s))); }1403
/** `s + a` reusing the expiring array @p a in place (commutative). @param s the arithmetic scalar (converted to @p T). @param a expiring array operand. @return the result, in @p a's buffer. @complexity O(size of @p a). @alloc none on the fast path. */1404
template <typename T, typename S>1405
requires std::is_arithmetic_v<S>1406
basic_ndarray<T> operator+(S s, basic_ndarray<T>&& a) { return add(std::move(a), scalar(static_cast<T>(s))); }1407
/** `a - s` reusing the expiring array @p a in place. @param a expiring array operand. @param s the arithmetic scalar (converted to @p T). @return the result, in @p a's buffer. @complexity O(size of @p a). @alloc none on the fast path. */1408
template <typename T, typename S>1409
requires std::is_arithmetic_v<S>1410
basic_ndarray<T> operator-(basic_ndarray<T>&& a, S s) { return sub(std::move(a), scalar(static_cast<T>(s))); }1411
/** `a * s` reusing the expiring array @p a in place. @param a expiring array operand. @param s the arithmetic scalar (converted to @p T). @return the result, in @p a's buffer. @complexity O(size of @p a). @alloc none on the fast path. */1412
template <typename T, typename S>1413
requires std::is_arithmetic_v<S>1414
basic_ndarray<T> operator*(basic_ndarray<T>&& a, S s) { return mul(std::move(a), scalar(static_cast<T>(s))); }1415
/** `s * a` reusing the expiring array @p a in place (commutative). @param s the arithmetic scalar (converted to @p T). @param a expiring array operand. @return the result, in @p a's buffer. @complexity O(size of @p a). @alloc none on the fast path. */1416
template <typename T, typename S>1417
requires std::is_arithmetic_v<S>1418
basic_ndarray<T> operator*(S s, basic_ndarray<T>&& a) { return mul(std::move(a), scalar(static_cast<T>(s))); }1419
/** `a / s` reusing the expiring array @p a in place. @param a expiring array operand. @param s the arithmetic scalar divisor (converted to @p T). @return the result, in @p a's buffer. @complexity O(size of @p a). @alloc none on the fast path. */1420
template <typename T, typename S>1421
requires std::is_arithmetic_v<S>1422
basic_ndarray<T> operator/(basic_ndarray<T>&& a, S s) { return divide(std::move(a), scalar(static_cast<T>(s))); }1423
/// @endcond1425
/// In-place compound assignment: `a += b`, `a -= b`, `a *= b`, `a /= b`1426
/// (array or arithmetic-scalar right operand). See compound_apply.1427
/**1428
* In-place `a += b`, updating @p a's buffer (see compound_apply).1429
* @param a the array to update in place. @param b the right operand (same shape or single-element).1430
* @return reference to @p a. @complexity O(size of @p a). @alloc none on the contiguous fast path.1431
* @test CheatahNDArray.CompoundAssignInPlace1432
* @test CheatahNDArray.CompoundAssignNonContiguousFallback1433
*/1434
template <typename T>1435
basic_ndarray<T>& operator+=(basic_ndarray<T>& a, const basic_ndarray<T>& b) {1436
compound_apply(a, b, detail::add_op{}); return a;1437
}1438
/**1439
* In-place `a -= b`, updating @p a's buffer (see compound_apply).1440
* @param a the array to update in place. @param b the right operand (same shape or single-element).1441
* @return reference to @p a. @complexity O(size of @p a). @alloc none on the contiguous fast path.1442
* @test CheatahNDArray.CompoundAssignNonContiguousFallback1443
*/1444
template <typename T>1445
basic_ndarray<T>& operator-=(basic_ndarray<T>& a, const basic_ndarray<T>& b) {1446
compound_apply(a, b, detail::sub_op{}); return a;1447
}1448
/**1449
* In-place `a *= b`, updating @p a's buffer (see compound_apply).1450
* @param a the array to update in place. @param b the right operand (same shape or single-element).1451
* @return reference to @p a. @complexity O(size of @p a). @alloc none on the contiguous fast path.1452
* @test CheatahNDArray.CompoundAssignNonContiguousFallback1453
*/1454
template <typename T>1455
basic_ndarray<T>& operator*=(basic_ndarray<T>& a, const basic_ndarray<T>& b) {1456
compound_apply(a, b, detail::mul_op{}); return a;1457
}1458
/**1459
* In-place `a /= b`, updating @p a's buffer (see compound_apply).1460
* @param a the array to update in place. @param b the right operand (same shape or single-element).1461
* @return reference to @p a. @complexity O(size of @p a). @alloc none on the contiguous fast path.1462
* @test CheatahNDArray.CompoundAssignNonContiguousFallback1463
*/1464
template <typename T>1465
basic_ndarray<T>& operator/=(basic_ndarray<T>& a, const basic_ndarray<T>& b) {1466
compound_apply(a, b, detail::div_op{}); return a;1467
}1468
/** In-place `a += s` (scalar). @param a the array to update. @param s the arithmetic scalar (converted to @p T). @return reference to @p a. @complexity O(size of @p a). @alloc none on the fast path. */1469
template <typename T, typename S>1470
requires std::is_arithmetic_v<S>1471
basic_ndarray<T>& operator+=(basic_ndarray<T>& a, S s) { return a += scalar(static_cast<T>(s)); }1472
/** In-place `a -= s` (scalar). @param a the array to update. @param s the arithmetic scalar (converted to @p T). @return reference to @p a. @complexity O(size of @p a). @alloc none on the fast path. @test CheatahNDArray.CompoundAssignInPlace */1473
template <typename T, typename S>1474
requires std::is_arithmetic_v<S>1475
basic_ndarray<T>& operator-=(basic_ndarray<T>& a, S s) { return a -= scalar(static_cast<T>(s)); }1476
/** In-place `a *= s` (scalar). @param a the array to update. @param s the arithmetic scalar (converted to @p T). @return reference to @p a. @complexity O(size of @p a). @alloc none on the fast path. @test CheatahNDArray.CompoundAssignInPlace */1477
template <typename T, typename S>1478
requires std::is_arithmetic_v<S>1479
basic_ndarray<T>& operator*=(basic_ndarray<T>& a, S s) { return a *= scalar(static_cast<T>(s)); }1480
/** In-place `a /= s` (scalar). @param a the array to update. @param s the arithmetic scalar divisor (converted to @p T). @return reference to @p a. @complexity O(size of @p a). @alloc none on the fast path. @test CheatahNDArray.CompoundAssignInPlace */1481
template <typename T, typename S>1482
requires std::is_arithmetic_v<S>1483
basic_ndarray<T>& operator/=(basic_ndarray<T>& a, S s) { return a /= scalar(static_cast<T>(s)); }1485
// ---- complex support ----1486
namespace detail {1487
/// Map @p a element-wise through @p f into a fresh contiguous array of element type1488
/// `U` (which may differ from `T` — e.g. complex→real for @ref real). Contiguous1489
/// fast path via `std::transform(unseq)`; otherwise a C-order walk.1490
template <typename U, Field T, typename F>1491
basic_ndarray<U> map_array(const basic_ndarray<T>& a, F f) {1492
basic_ndarray<U> out = basic_ndarray<U>::uninitialized(a.shape()); // fully written below1493
auto& obuf = *out.buffer();1494
if (is_contiguous(a)) {1495
const auto& abuf = *a.buffer();1496
std::transform(CHEATAH_UNSEQ abuf.begin() + a.offset(),1497
abuf.begin() + a.offset() + a.size(), obuf.begin(), f);1498
return out;1499
}1500
std::vector<std::size_t> idx(a.ndim(), 0);1501
std::size_t flat = 0;1502
do {1503
obuf[flat++] = f(a.at(idx));1504
} while (a.ndim() != 0 && next_index(idx, a.shape()));1505
return out;1506
}1508
// Out-of-line, separately-compiled (-ffast-math) double-precision SIMD kernels for the1509
// element-wise ufuncs — see ufunc_simd.cpp. They vectorize the transcendentals through1510
// libmvec, which the default flags cannot; isolating -ffast-math to that file keeps the1511
// rest of cheatah's arithmetic strict.1512
void simd_sqrt_f64(const double*, double*, std::size_t);1513
void simd_cbrt_f64(const double*, double*, std::size_t);1514
void simd_exp_f64(const double*, double*, std::size_t);1515
void simd_log_f64(const double*, double*, std::size_t);1516
void simd_sin_f64(const double*, double*, std::size_t);1517
void simd_cos_f64(const double*, double*, std::size_t);1518
void simd_tan_f64(const double*, double*, std::size_t);1520
/// Map a ufunc over @p a: a *contiguous double* array goes through the precompiled SIMD1521
/// @p kernel; everything else (float, or a strided/broadcast view) uses the generic1522
/// scalar @p fallback. Same result either way — the kernel just vectorizes the hot case.1523
template <FloatingPoint T, class Kernel, class Fallback>1524
basic_ndarray<T> map_ufunc(const basic_ndarray<T>& a, Kernel kernel, Fallback fallback) {1525
if constexpr (std::is_same_v<T, double>) {1526
if (is_contiguous(a)) {1527
basic_ndarray<T> out = basic_ndarray<T>::uninitialized(a.shape()); // kernel fills it1528
kernel(a.buffer()->data() + a.offset(), out.buffer()->data(), a.size());1529
return out;1530
}1531
}1532
return map_array<T>(a, fallback);1533
}1534
} // namespace detail1536
/**1537
* Build a complex array from real and imaginary parts (element-wise `re + im·j`),1538
* broadcasting the two together — the way to construct a complex matrix/vector1539
* (a wavefunction, a Hermitian operator) since cheatah literals are real.1540
* @param re the real parts (a real floating array).1541
* @param im the imaginary parts (a real floating array, broadcast against @p re).1542
* @return a `basic_ndarray<std::complex<T>>` of `re + im·j`; throws if the shapes don't broadcast.1543
* @complexity O(size of result).1544
* @alloc allocates the result buffer.1545
* @test CheatahNDArray.ComplexConstructAndParts1546
* @crtest NdarrayCompileRun.Complex1547
* @systest StdlibE2E.NdarrayComplex1548
*/1549
template <FloatingPoint T>1550
basic_ndarray<std::complex<T>> complex(const basic_ndarray<T>& re, const basic_ndarray<T>& im) {1551
using C = std::complex<T>;1552
const std::vector<std::size_t> rshape = broadcast_shapes(re.shape(), im.shape());1553
const basic_ndarray<T> rv = broadcast_to(re, rshape);1554
const basic_ndarray<T> iv = broadcast_to(im, rshape);1555
basic_ndarray<C> out(rshape);1556
auto& obuf = *out.buffer();1557
std::vector<std::size_t> idx(rshape.size(), 0);1558
std::size_t flat = 0;1559
do {1560
obuf[flat++] = C(rv.at(idx), iv.at(idx));1561
} while (!rshape.empty() && detail::next_index(idx, rshape));1562
return out;1563
}1564
/**1565
* Element-wise complex conjugate (`a − b·j` for each `a + b·j`); on a real array it1566
* is the identity (a copy). Type-preserving. Used to form Hermitian adjoints and1567
* conjugate-linear inner products.1568
* @param a the array.1569
* @return a fresh array of the same element type with each element conjugated.1570
* @complexity O(size).1571
* @alloc allocates the result buffer.1572
* @test CheatahNDArray.ComplexConstructAndParts1573
* @crtest NdarrayCompileRun.Conj1574
* @systest StdlibE2E.NdarrayComplex1575
*/1576
template <Field T>1577
basic_ndarray<T> conj(const basic_ndarray<T>& a) {1578
return detail::map_array<T>(a, [](T x) -> T {1579
if constexpr (is_complex_v<T>) {1580
return std::conj(x);1581
} else {1582
return x;1583
}1584
});1585
}1586
/**1587
* The real parts as a real array (the identity on a real array). For `a + b·j` it1588
* returns `a`.1589
* @param a the array.1590
* @return a `basic_ndarray<real_base_t<T>>` of the real parts.1591
* @complexity O(size).1592
* @alloc allocates the result buffer.1593
* @test CheatahNDArray.ComplexConstructAndParts1594
* @crtest NdarrayCompileRun.Real1595
* @systest StdlibE2E.NdarrayComplex1596
*/1597
template <Field T>1598
basic_ndarray<real_base_t<T>> real(const basic_ndarray<T>& a) {1599
using R = real_base_t<T>;1600
return detail::map_array<R>(a, [](T x) -> R {1601
if constexpr (is_complex_v<T>) {1602
return x.real();1603
} else {1604
return x;1605
}1606
});1607
}1608
/**1609
* The imaginary parts as a real array (all zeros for a real array). For `a + b·j` it1610
* returns `b`.1611
* @param a the array.1612
* @return a `basic_ndarray<real_base_t<T>>` of the imaginary parts.1613
* @complexity O(size).1614
* @alloc allocates the result buffer.1615
* @test CheatahNDArray.ComplexConstructAndParts1616
* @crtest NdarrayCompileRun.Imag1617
* @systest StdlibE2E.NdarrayComplex1618
*/1619
template <Field T>1620
basic_ndarray<real_base_t<T>> imag(const basic_ndarray<T>& a) {1621
using R = real_base_t<T>;1622
return detail::map_array<R>(a, [](T x) -> R {1623
if constexpr (is_complex_v<T>) {1624
return x.imag();1625
} else {1626
return R{0};1627
}1628
});1629
}1631
// ---- element-wise math (numpy-style ufuncs) ----1632
// These are the array counterparts of the scalar `math` module — mirroring Python's1633
// split: `math.sqrt(x)` for a scalar, `ndarray.sqrt(a)` (≈ `numpy.sqrt`) for a whole1634
// array. A contiguous `double` array routes through a precompiled SIMD kernel1635
// (ufunc_simd.cpp) that vectorizes via glibc's libmvec — so `exp`/`sin`/… run at vector1636
// speed and beat NumPy's ufuncs; other element types / strided views fall back to a1637
// scalar map (see detail::map_ufunc).1638
/**1639
* Element-wise square root (the array form of `math.sqrt`; ≈ `numpy.sqrt`).1640
* @param a a floating-point array.1641
* @return a fresh same-shape array with `√x` for each element.1642
* @complexity O(size). @alloc allocates the result buffer.1643
* @test CheatahNDArray.ElementwiseMath1644
* @crtest NdarrayCompileRun.Sqrt1645
* @systest StdlibE2E.NdarrayMath1646
*/1647
template <FloatingPoint T>1648
basic_ndarray<T> sqrt(const basic_ndarray<T>& a) {1649
return detail::map_ufunc<T>(a, detail::simd_sqrt_f64, [](T x) { return std::sqrt(x); });1650
}1651
/**1652
* Element-wise cube root (the array form of `math.cbrt`; ≈ `numpy.cbrt`).1653
* @param a a floating-point array.1654
* @return a fresh same-shape array with `∛x` for each element.1655
* @complexity O(size). @alloc allocates the result buffer.1656
* @test CheatahNDArray.ElementwiseMath1657
* @systest StdlibE2E.NdarrayMath1658
*/1659
template <FloatingPoint T>1660
basic_ndarray<T> cbrt(const basic_ndarray<T>& a) {1661
return detail::map_ufunc<T>(a, detail::simd_cbrt_f64, [](T x) { return std::cbrt(x); });1662
}1663
/**1664
* Element-wise eˣ (the array form of `math.exp`; ≈ `numpy.exp`).1665
* @param a a floating-point array.1666
* @return a fresh same-shape array with `exp(x)` for each element.1667
* @complexity O(size). @alloc allocates the result buffer.1668
* @test CheatahNDArray.ElementwiseMath1669
* @crtest NdarrayCompileRun.Exp1670
* @systest StdlibE2E.NdarrayMath1671
*/1672
template <FloatingPoint T>1673
basic_ndarray<T> exp(const basic_ndarray<T>& a) {1674
return detail::map_ufunc<T>(a, detail::simd_exp_f64, [](T x) { return std::exp(x); });1675
}1676
/**1677
* Element-wise natural log (the array form of `math.log`; ≈ `numpy.log`).1678
* @param a a floating-point array.1679
* @return a fresh same-shape array with `ln(x)` for each element.1680
* @complexity O(size). @alloc allocates the result buffer.1681
* @test CheatahNDArray.ElementwiseMath1682
* @systest StdlibE2E.NdarrayMath1683
*/1684
template <FloatingPoint T>1685
basic_ndarray<T> log(const basic_ndarray<T>& a) {1686
return detail::map_ufunc<T>(a, detail::simd_log_f64, [](T x) { return std::log(x); });1687
}1688
/**1689
* Element-wise sine (the array form of `math.sin`; ≈ `numpy.sin`).1690
* @param a a floating-point array (radians).1691
* @return a fresh same-shape array with `sin(x)` for each element.1692
* @complexity O(size). @alloc allocates the result buffer.1693
* @test CheatahNDArray.ElementwiseMath1694
* @crtest NdarrayCompileRun.Sin1695
* @systest StdlibE2E.NdarrayMath1696
*/1697
template <FloatingPoint T>1698
basic_ndarray<T> sin(const basic_ndarray<T>& a) {1699
return detail::map_ufunc<T>(a, detail::simd_sin_f64, [](T x) { return std::sin(x); });1700
}1701
/**1702
* Element-wise cosine (the array form of `math.cos`; ≈ `numpy.cos`).1703
* @param a a floating-point array (radians).1704
* @return a fresh same-shape array with `cos(x)` for each element.1705
* @complexity O(size). @alloc allocates the result buffer.1706
* @test CheatahNDArray.ElementwiseMath1707
* @systest StdlibE2E.NdarrayMath1708
*/1709
template <FloatingPoint T>1710
basic_ndarray<T> cos(const basic_ndarray<T>& a) {1711
return detail::map_ufunc<T>(a, detail::simd_cos_f64, [](T x) { return std::cos(x); });1712
}1713
/**1714
* Element-wise tangent (the array form of `math.tan`; ≈ `numpy.tan`).1715
* @param a a floating-point array (radians).1716
* @return a fresh same-shape array with `tan(x)` for each element.1717
* @complexity O(size). @alloc allocates the result buffer.1718
* @test CheatahNDArray.ElementwiseMath1719
* @systest StdlibE2E.NdarrayMath1720
*/1721
template <FloatingPoint T>1722
basic_ndarray<T> tan(const basic_ndarray<T>& a) {1723
return detail::map_ufunc<T>(a, detail::simd_tan_f64, [](T x) { return std::tan(x); });1724
}1725
/**1726
* Element-wise absolute value (the array form of `math.abs`; ≈ `numpy.abs`).1727
* @param a a floating-point array.1728
* @return a fresh same-shape array with `|x|` for each element.1729
* @complexity O(size). @alloc allocates the result buffer.1730
* @test CheatahNDArray.ElementwiseMath1731
* @systest StdlibE2E.NdarrayMath1732
*/1733
template <FloatingPoint T>1734
basic_ndarray<T> abs(const basic_ndarray<T>& a) {1735
return detail::map_array<T>(a, [](T x) { return std::fabs(x); });1736
}1738
// ---- reductions / access / display ----1739
namespace detail {1740
/// The shared multi-accumulator reduction: sums `get(0)..get(n-1)` with EIGHT independent1741
/// accumulators, tree-combined, plus a scalar tail. The independent lanes break the FP-add1742
/// dependency chain so -O3 -march=native emits SIMD+FMA and reaches memory bandwidth instead of1743
/// add latency (a single running sum — or a plain `std::reduce`, which libstdc++ left-folds for FP1744
/// without -ffast-math — serializes: the dot/norm mistake). `get(i)` returns the i-th TERM — an1745
/// element for `sum`, a (possibly conjugated) product for `dot`, a strided read for `trace`. One1746
/// primitive replaces the copies formerly hand-rolled in ndarray/linalg. `constexpr`, so a1747
/// fixed-extent caller gets a compile-time reduction too.1748
template <class T, class Get>1749
constexpr T reduce_lanes(std::size_t n, Get get) {1750
T s0{}, s1{}, s2{}, s3{}, s4{}, s5{}, s6{}, s7{};1751
std::size_t i = 0;1752
for (; i + 8 <= n; i += 8) {1753
s0 += get(i + 0); s1 += get(i + 1); s2 += get(i + 2); s3 += get(i + 3);1754
s4 += get(i + 4); s5 += get(i + 5); s6 += get(i + 6); s7 += get(i + 7);1755
}1756
T s = ((s0 + s1) + (s2 + s3)) + ((s4 + s5) + (s6 + s7));1757
for (; i < n; ++i) s += get(i);1758
return s;1759
}1760
} // namespace detail1761
/**1762
* Sum of all elements — a full reduction across every axis (a contiguous array goes1763
* through the shared multi-accumulator SIMD reduction @ref detail::reduce_lanes, else1764
* a C-order walk); empty sums to 0.1765
* @param a the array.1766
* @return the total, as the element type @p T.1767
* @complexity O(size).1768
* @alloc none.1769
* @test CheatahNDArray.ShapeFactoriesAndReductions1770
* @crtest NdarrayCompileRun.Sum1771
* @systest StdlibE2E.Ndarray1772
*/1773
template <Field T>1774
T sum(const basic_ndarray<T>& a) {1775
if (is_contiguous(a)) {1776
// Contiguous fast path via the shared multi-accumulator reduction (each term is one element).1777
const T* p = a.buffer()->data() + a.offset();1778
return detail::reduce_lanes<T>(a.size(), [p](std::size_t i) { return p[i]; });1779
}1780
T s{};1781
std::vector<std::size_t> idx(a.ndim(), 0);1782
do {1783
s += a.at(idx);1784
} while (a.ndim() != 0 && detail::next_index(idx, a.shape()));1785
return s;1786
}1787
/**1788
* Mean of all elements, always as a `double` (0.0 for an empty array — no divide-by-zero).1789
* @param a the array.1790
* @return the average as a double.1791
* @complexity O(size).1792
* @alloc none.1793
* @test CheatahNDArray.ShapeFactoriesAndReductions1794
* @crtest NdarrayCompileRun.Mean1795
* @systest StdlibE2E.Ndarray1796
*/1797
template <Numeric T>1798
double mean(const basic_ndarray<T>& a) {1799
const std::size_t n = a.size();1800
return n == 0 ? 0.0 : static_cast<double>(sum(a)) / static_cast<double>(n);1801
}1802
/**1803
* Read one element by signed multi-index (the cheatah-facing wrapper over @ref1804
* basic_ndarray::at; rejects negative coordinates).1805
* @param a the array.1806
* @param index one coordinate per dimension (signed; throws on a negative).1807
* @return the element value (type @p T); throws on a wrong-rank/out-of-range index.1808
* @complexity O(ndim).1809
* @alloc none.1810
* @test CheatahNDArray.ShapeFactoriesAndReductions1811
* @crtest NdarrayCompileRun.Get1812
* @systest StdlibE2E.Ndarray1813
*/1814
template <Copyable T>1815
T get(const basic_ndarray<T>& a, const std::vector<long long>& index) {1816
return a.at(detail::to_size(index));1817
}1818
/**1819
* The shape as signed dims (cheatah integers are signed; a 0-d array yields an empty list).1820
* @param a the array.1821
* @return the dimensions as a `long long` vector.1822
* @complexity O(ndim).1823
* @alloc allocates the result vector.1824
* @test CheatahNDArray.ShapeFactoriesAndReductions1825
* @crtest NdarrayCompileRun.ShapeOf1826
* @systest StdlibE2E.Ndarray1827
*/1828
template <Element T>1829
std::vector<long long> shape_of(const basic_ndarray<T>& a) {1830
std::vector<long long> out(a.ndim());1831
for (std::size_t i = 0; i < a.ndim(); ++i) out[i] = static_cast<long long>(a.shape()[i]);1832
return out;1833
}1834
/**1835
* The element count as a signed value (1 for a 0-d array).1836
* @param a the array.1837
* @return the number of elements as a `long long`.1838
* @complexity O(ndim).1839
* @alloc none.1840
* @test CheatahNDArray.ShapeFactoriesAndReductions1841
* @crtest NdarrayCompileRun.SizeOf1842
* @systest StdlibE2E.Ndarray1843
*/1844
template <Element T>1845
long long size_of(const basic_ndarray<T>& a) {1846
return static_cast<long long>(a.size());1847
}1849
namespace detail {1850
/// Format one element. Real types go through `operator<<`; a complex element is1851
/// rendered Python-style as `a+bj` / `a-bj` (not the `std::complex` default1852
/// `(a,b)`), so a complex spectrum prints the way a cheatah user expects.1853
template <typename T>1854
std::string format_scalar(const T& v) {1855
std::ostringstream os;1856
if constexpr (is_complex_v<T>) {1857
using R = real_base_t<T>;1858
// Flush negative zero to +0 so a conjugate prints "1+0j", not "1+-0j".1859
const auto nz = [](R x) -> R { return x == R{0} ? R{0} : x; };1860
os << nz(v.real());1861
if (v.imag() < R{0}) {1862
os << "-" << nz(-v.imag()) << "j";1863
} else {1864
os << "+" << nz(v.imag()) << "j";1865
}1866
} else if constexpr (std::is_same_v<T, signed char> || std::is_same_v<T, unsigned char> ||1867
std::is_same_v<T, char>) {1868
os << +v; // i8/u8 are char-sized: promote so an element prints as a NUMBER, not a character1869
} else {1870
os << v;1871
}1872
return os.str();1873
}1875
/// Recursively format @p a into nested brackets (each element via `format_scalar`).1876
template <Element T>1877
void format_rec(const basic_ndarray<T>& a, std::vector<std::size_t>& idx, std::size_t dim,1878
std::string& out) {1879
if (dim == a.ndim()) {1880
out += format_scalar(a.at(idx));1881
return;1882
}1883
out += "[";1884
for (std::size_t i = 0; i < a.shape()[dim]; ++i) {1885
if (i != 0) out += ", ";1886
idx[dim] = i;1887
format_rec(a, idx, dim + 1, out);1888
}1889
out += "]";1890
}1892
/// Like format_rec but ABBREVIATES a large array: an axis longer than `2*edge` shows its1893
/// first and last `edge` items with `...` between, recursively. Summarization is enabled by1894
/// @p summarize (the caller turns it on only past a total-size threshold), so small arrays1895
/// print in full.1896
template <Element T>1897
void format_rec_trunc(const basic_ndarray<T>& a, std::vector<std::size_t>& idx, std::size_t dim,1898
std::string& out, std::size_t edge, bool summarize) {1899
if (dim == a.ndim()) {1900
out += format_scalar(a.at(idx));1901
return;1902
}1903
const std::size_t n = a.shape()[dim];1904
const bool trunc = summarize && n > 2 * edge;1905
out += "[";1906
bool first = true;1907
for (std::size_t i = 0; i < n; ++i) {1908
if (trunc && i >= edge && i < n - edge) {1909
if (i == edge) {1910
if (!first) out += ", ";1911
out += "...";1912
first = false;1913
}1914
continue; // skip the abbreviated middle1915
}1916
if (!first) out += ", ";1917
first = false;1918
idx[dim] = i;1919
format_rec_trunc(a, idx, dim + 1, out, edge, summarize);1920
}1921
out += "]";1922
}1924
/// to_string, but ABBREVIATED with `...` when the array is large (total size beyond a1925
/// numpy-style threshold) — the readable default for `io.print`. `io.rprint`/`to_string`1926
/// keep the full form.1927
template <Element T>1928
std::string to_string_pretty(const basic_ndarray<T>& a) {1929
if (a.ndim() == 0) return format_scalar(a.at({}));1930
constexpr std::size_t kEdge = 3; // items kept at each end of an abbreviated axis1931
constexpr std::size_t kThreshold = 1000; // only summarize past this many elements (numpy)1932
std::vector<std::size_t> idx(a.ndim(), 0);1933
std::string out;1934
format_rec_trunc(a, idx, 0, out, kEdge, a.size() > kThreshold);1935
return out;1936
}1937
} // namespace detail1939
/**1940
* Render as a nested-bracket string, e.g. `"[[1, 2], [3, 4]]"` (a 0-d scalar renders1941
* as the bare number). Each element is formatted with the default `ostream` precision.1942
* @param a the array.1943
* @return the textual representation.1944
* @complexity O(size).1945
* @alloc allocates the result string.1946
* @test CheatahNDArray.ToStringScalar, CheatahNDArray.BroadcastingAdd1947
* @crtest NdarrayCompileRun.ToString1948
* @systest StdlibE2E.Ndarray1949
*/1950
template <Element T>1951
std::string to_string(const basic_ndarray<T>& a) {1952
if (a.ndim() == 0) {1953
return detail::format_scalar(a.at({}));1954
}1955
std::vector<std::size_t> idx(a.ndim(), 0);1956
std::string out;1957
detail::format_rec(a, idx, 0, out);1958
return out;1959
}1961
template <Element T>1962
inline std::string basic_ndarray<T>::str() const {1963
return to_string(*this);1964
}1966
/**1967
* Stream an array to a `std::ostream` (the FULL nested-bracket form) — so an NDArray is1968
* directly Streamable, like a primitive or a cheatah struct, without going through1969
* `to_string`/`str()`. `io.rprint`, `str()`, and a struct that holds an array all stream it1970
* this way; `io.print` instead uses @ref cheatah_pretty_print to abbreviate large arrays.1971
* @param os the stream.1972
* @param a the array.1973
* @return @p os.1974
* @complexity O(size).1975
* @alloc allocates the intermediate string.1976
* @test CheatahNDArray.StreamableOperator1977
* @systest StdlibE2E.Ndarray1978
*/1979
template <Element T>1980
std::ostream& operator<<(std::ostream& os, const basic_ndarray<T>& a) {1981
return os << to_string(a);1982
}1984
template <Element T>1985
inline void basic_ndarray<T>::cheatah_pretty_print(std::ostream& os, long long) const {1986
os << detail::to_string_pretty(*this);1987
}1989
} // namespace cheatah::ndarray1991
// cheatah's value-position subscript lowers to builtins::index(obj, i, ...).1992
// These overloads give it the ndarray meaning: negative-aware element reads,1993
// one coordinate per dimension.1994
namespace cheatah::builtins {1996
/**1997
* Element read `a[i, j, ...]` (negative indices count from the dimension end).1998
* @param a the array to read from.1999
* @param first the first-axis coordinate (negative counts from that dimension's end).2000
* @param rest the remaining per-axis coordinates (one per further dimension).2001
* @return the element value.2002
* @complexity O(ndim). @alloc none.2003
* @test CheatahNDArray.SubscriptReadWrite2004
* @crtest LangFeatures.NdarraySubscript2005
* @systest StdlibE2E.Ndarray2006
*/2007
template <typename T, ::cheatah::ndarray::Subscript First, ::cheatah::ndarray::Subscript... Ix>2008
T index(const ::cheatah::ndarray::basic_ndarray<T>& a, First first, Ix... rest) {2009
// const_cast is sound: item_ref only computes a position; we copy the value out. Any index may be2010
// a scoped enum column label — item_ref performs the one sanctioned conversion (ndarray::Subscript).2011
return const_cast<::cheatah::ndarray::basic_ndarray<T>&>(a).item_ref(first, rest...);2012
}2014
} // namespace cheatah::builtins