Source
stdlib/ndarray/ndarray.cpp
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
#include "ndarray.hpp"5
#include <algorithm>6
#include <stdexcept>8
namespace cheatah::ndarray {10
// Every other ndarray function is now a template over the element type and lives11
// in the header (ndarray.hpp). `broadcast_shapes` is the one shape-only,12
// element-type-independent function, so it is compiled here — giving the library a13
// translation unit and a real symbol while the templated ops monomorphize at the14
// call site.15
std::vector<std::size_t> broadcast_shapes(const std::vector<std::size_t>& a,16
const std::vector<std::size_t>& b) {17
const std::size_t n = std::max(a.size(), b.size());18
std::vector<std::size_t> r(n);19
for (std::size_t i = 0; i < n; ++i) {20
// Align from the right (trailing dimensions).21
const std::size_t da = (i < n - a.size()) ? 1 : a[i - (n - a.size())];22
const std::size_t db = (i < n - b.size()) ? 1 : b[i - (n - b.size())];23
if (da == db || db == 1) {24
r[i] = da;25
} else if (da == 1) {26
r[i] = db;27
} else {28
throw std::runtime_error("ndarray: operands could not be broadcast together");29
}30
}31
return r;32
}34
} // namespace cheatah::ndarray