From 03e23202e2d32195d1f5112f51fe499cd970a079 Mon Sep 17 00:00:00 2001 From: Jacob Domagala Date: Wed, 12 Nov 2025 13:23:45 +0100 Subject: [PATCH 1/9] [#80]: Add .codex to gitignore --- .gitignore | 1 + 1 file changed, 1 insertion(+) diff --git a/.gitignore b/.gitignore index d7a8a80..af6abaa 100644 --- a/.gitignore +++ b/.gitignore @@ -10,3 +10,4 @@ __pycache__/ *.pyo venv/ .env/ +.codex From 1c2ef11f4e918ff1d7d06bd6e35426271422f259 Mon Sep 17 00:00:00 2001 From: Jacob Domagala Date: Wed, 12 Nov 2025 12:00:00 +0100 Subject: [PATCH 2/9] [#80]: Add rank-8 tensor shape metadata --- src/autograd.hpp | 6 +++--- src/tensor.hpp | 24 ++++++++++++++---------- src/types.hpp | 4 ++++ 3 files changed, 21 insertions(+), 13 deletions(-) diff --git a/src/autograd.hpp b/src/autograd.hpp index 15d8480..80c45e7 100644 --- a/src/autograd.hpp +++ b/src/autograd.hpp @@ -63,8 +63,8 @@ struct MatmulFunction : Function { void print() const override { printf("MatmulFunction\n"); } private: - std::array lhs_shape; - std::array rhs_shape; + TensorShape lhs_shape; + TensorShape rhs_shape; }; struct ReLUFunction : Function { @@ -103,7 +103,7 @@ struct SumFunction : Function { private: int64_t dim_; - std::array input_shape_; + TensorShape input_shape_; }; struct MseFunction : Function { diff --git a/src/tensor.hpp b/src/tensor.hpp index f95df1b..ab86bdd 100644 --- a/src/tensor.hpp +++ b/src/tensor.hpp @@ -2,9 +2,9 @@ #include "types.hpp" -#include #include #include +#include namespace smollnet { @@ -25,8 +25,8 @@ struct Storage { struct TensorImpl { std::shared_ptr storage = nullptr; - std::array sizes = {0, 0, 0}; - std::array strides = {0, 0, 0}; + TensorShape sizes = {}; + TensorShape strides = {}; bool expanded = false; size_t elems = 1; @@ -76,8 +76,8 @@ class Tensor { DataType dtype() const noexcept; void *data() const noexcept; size_t numel() const noexcept; - const std::array& dims() const noexcept; - const std::array& strides() const noexcept; + const TensorShape &dims() const noexcept; + const TensorShape &strides() const noexcept; void print() const; void print_elms() const; std::string to_string() const; @@ -98,7 +98,7 @@ class Tensor { Tensor matmul(const Tensor&other) const; Tensor transpose(int d0, int d1) const; - Tensor expand(const std::array &new_sz) const; + Tensor expand(const TensorShape &new_sz) const; Tensor cuda() const; Tensor cpu() const; @@ -165,28 +165,32 @@ void manual_seed(unsigned long long seed); template Tensor empty(const int64_t (&dims)[N], DataType t, Device d, bool requires_grad = false) { - static_assert(N <= 3, "We don't support more than 3 dimensional Tensors"); + static_assert(N <= kMaxTensorDims, + "We don't support more than kMaxTensorDims dimensional Tensors"); return empty(dims, N, t, d, requires_grad); } template Tensor zeros(const int64_t (&dims)[N], DataType t, Device d, bool requires_grad = false) { - static_assert(N <= 3, "We don't support more than 3 dimensional Tensors"); + static_assert(N <= kMaxTensorDims, + "We don't support more than kMaxTensorDims dimensional Tensors"); return zeros(dims, N, t, d, requires_grad); } template Tensor ones(const int64_t (&dims)[N], DataType t, Device d, bool requires_grad = false) { - static_assert(N <= 3, "We don't support more than 3 dimensional Tensors"); + static_assert(N <= kMaxTensorDims, + "We don't support more than kMaxTensorDims dimensional Tensors"); return ones(dims, N, t, d, requires_grad); } template Tensor rand(const int64_t (&dims)[N], DataType t, Device d, bool requires_grad = false) { - static_assert(N <= 3, "We don't support more than 3 dimensional Tensors"); + static_assert(N <= kMaxTensorDims, + "We don't support more than kMaxTensorDims dimensional Tensors"); return rand(dims, N, t, d, requires_grad); } diff --git a/src/types.hpp b/src/types.hpp index eca1655..60aeceb 100644 --- a/src/types.hpp +++ b/src/types.hpp @@ -2,9 +2,13 @@ #include #include +#include namespace smollnet { +inline constexpr size_t kMaxTensorDims = 8; +using TensorShape = std::array; + enum class Device : uint8_t { CUDA, CPU }; enum class DataType : uint8_t { From 75f4b9c3c35c268660568f1558b47ac035bfe3bf Mon Sep 17 00:00:00 2001 From: Jacob Domagala Date: Wed, 12 Nov 2025 12:01:00 +0100 Subject: [PATCH 3/9] [#80]: Generalize CUDA shape metadata --- src/kernels.cuh | 21 ++++++++------- src/operators.cu | 70 +++++++++++++++++------------------------------- 2 files changed, 36 insertions(+), 55 deletions(-) diff --git a/src/kernels.cuh b/src/kernels.cuh index c94a96a..b19b3a6 100644 --- a/src/kernels.cuh +++ b/src/kernels.cuh @@ -4,6 +4,8 @@ #include #include +#include "types.hpp" + namespace smollnet { constexpr int32_t ROW_MAJOR = 0; @@ -12,24 +14,24 @@ constexpr int32_t DEPTH_MAJOR = 2; struct StrideAndSize { - std::array stride; + int64_t stride[kMaxTensorDims] = {}; int64_t rank; - std::array size; + int64_t size[kMaxTensorDims] = {}; }; struct StrideInfo { // size of the output operation - int64_t output_size[3]; + int64_t output_size[kMaxTensorDims] = {}; - int64_t a_stride[3]; - int64_t b_stride[3]; + int64_t a_stride[kMaxTensorDims] = {}; + int64_t b_stride[kMaxTensorDims] = {}; int64_t rank; }; struct SizeInfo { - int64_t a_size[3]; - int64_t b_size[3]; + int64_t a_size[kMaxTensorDims] = {}; + int64_t b_size[kMaxTensorDims] = {}; }; enum class WelfordType : uint8_t{ @@ -58,9 +60,8 @@ void launch_div(float *out, float *a, float *b, size_t numElems); void launch_div_strided(void *dst, void *a, void *b, const StrideInfo &s, size_t total); -void launch_sum_dim0(void *out, void *in, const StrideAndSize& s_input, const StrideAndSize& s_output); -void launch_sum_dim1(void *out, void *in, const StrideAndSize& s_input, const StrideAndSize& s_output); -void launch_sum_dim2(void *out, void *in, const StrideAndSize& s_input, const StrideAndSize& s_output); +void launch_sum_dim(void *out, void *in, const StrideAndSize &s_input, + const StrideAndSize &s_output, int64_t dim); void launch_matmul(void *out, void *left, void *right, const StrideInfo &strides, const SizeInfo &sizes, diff --git a/src/operators.cu b/src/operators.cu index d090b9a..b399c1d 100644 --- a/src/operators.cu +++ b/src/operators.cu @@ -5,23 +5,19 @@ namespace smollnet { -__device__ __forceinline__ void compute_dimensions(int (&dims)[3], size_t idx, - const StrideInfo &s) { - - if (s.rank == 3) { - int64_t rest = s.output_size[1] * s.output_size[2]; - dims[0] = idx / rest; - int64_t rem = idx % rest; - dims[1] = rem / s.output_size[2]; - dims[2] = rem % s.output_size[2]; - } else if (s.rank == 2) { - dims[0] = idx / s.output_size[1]; - dims[1] = idx % s.output_size[1]; - dims[2] = 0; - } else { // rank == 1 - dims[0] = idx; - dims[1] = 0; - dims[2] = 0; +__device__ __forceinline__ void compute_strided_offsets(size_t idx, + const StrideInfo &s, + int64_t &offA, + int64_t &offB) { + offA = 0; + offB = 0; + + for (int64_t dim = s.rank - 1; dim >= 0; --dim) { + const int64_t coord = idx % s.output_size[dim]; + idx /= s.output_size[dim]; + + offA += coord * s.a_stride[dim]; + offB += coord * s.b_stride[dim]; } } @@ -73,13 +69,9 @@ __global__ void add_strided_kernel(float *__restrict__ out, if (idx >= total) return; - int dims[3] = {0, 0, 0}; - compute_dimensions(dims, idx, s); - - int64_t offA = dims[0] * s.a_stride[0] + dims[1] * s.a_stride[1] + - dims[2] * s.a_stride[2]; - int64_t offB = dims[0] * s.b_stride[0] + dims[1] * s.b_stride[1] + - dims[2] * s.b_stride[2]; + int64_t offA = 0; + int64_t offB = 0; + compute_strided_offsets(idx, s, offA, offB); out[idx] = a[offA] + b[offB]; } @@ -124,13 +116,9 @@ __global__ void mul_strided_kernel(float *__restrict__ out, if (idx >= total) return; - int dims[3] = {0, 0, 0}; - compute_dimensions(dims, idx, s); - - int64_t offA = dims[0] * s.a_stride[0] + dims[1] * s.a_stride[1] + - dims[2] * s.a_stride[2]; - int64_t offB = dims[0] * s.b_stride[0] + dims[1] * s.b_stride[1] + - dims[2] * s.b_stride[2]; + int64_t offA = 0; + int64_t offB = 0; + compute_strided_offsets(idx, s, offA, offB); out[idx] = a[offA] * b[offB]; } @@ -167,13 +155,9 @@ __global__ void sub_strided_kernel(float *out, float *a, float *b, StrideInfo s, if (idx >= total) return; - int dims[3] = {0, 0, 0}; - compute_dimensions(dims, idx, s); - - int64_t offA = dims[0] * s.a_stride[0] + dims[1] * s.a_stride[1] + - dims[2] * s.a_stride[2]; - int64_t offB = dims[0] * s.b_stride[0] + dims[1] * s.b_stride[1] + - dims[2] * s.b_stride[2]; + int64_t offA = 0; + int64_t offB = 0; + compute_strided_offsets(idx, s, offA, offB); out[idx] = a[offA] - b[offB]; } @@ -210,13 +194,9 @@ __global__ void div_strided_kernel(float *__restrict__ out, if (idx >= total) return; - int dims[3] = {0, 0, 0}; - compute_dimensions(dims, idx, s); - - int64_t offA = dims[0] * s.a_stride[0] + dims[1] * s.a_stride[1] + - dims[2] * s.a_stride[2]; - int64_t offB = dims[0] * s.b_stride[0] + dims[1] * s.b_stride[1] + - dims[2] * s.b_stride[2]; + int64_t offA = 0; + int64_t offB = 0; + compute_strided_offsets(idx, s, offA, offB); out[idx] = a[offA] / b[offB]; } From 6d409e25ee09f68819573be8f615b0ad71bc6e8d Mon Sep 17 00:00:00 2001 From: Jacob Domagala Date: Wed, 12 Nov 2025 12:02:00 +0100 Subject: [PATCH 4/9] [#80]: Generalize tensor broadcasting --- src/tensor.cpp | 547 +++++++++++++++++++++++-------------------------- 1 file changed, 258 insertions(+), 289 deletions(-) diff --git a/src/tensor.cpp b/src/tensor.cpp index 17405be..b9be32a 100644 --- a/src/tensor.cpp +++ b/src/tensor.cpp @@ -6,7 +6,9 @@ #include #include #include +#include #include +#include #include @@ -23,6 +25,214 @@ std::mt19937 &cpu_random_generator() { return generator; } +std::string shape_to_string(const TensorShape &shape, int64_t rank) { + return fmt::format("[{}]", + fmt::join(shape.begin(), shape.begin() + rank, ", ")); +} + +int64_t infer_rank(const TensorShape &shape) { + int64_t rank = 0; + for (size_t dim = 0; dim < shape.size(); ++dim) { + if (shape[dim] != 0) { + rank = static_cast(dim + 1); + } + } + return rank; +} + +bool is_dense_contiguous(const Tensor &t) { + int64_t expected_stride = 1; + for (int64_t dim = t.ndims() - 1; dim >= 0; --dim) { + if (t.strides()[dim] != expected_stride) { + return false; + } + expected_stride *= t.size(dim); + } + return true; +} + +Tensor make_broadcast_view(const Tensor &input, const TensorShape &out_shape, + int64_t out_rank) { + ASSERT(out_rank >= input.ndims(), + fmt::format("Cannot broadcast tensor rank {} to rank {}", + input.ndims(), out_rank)); + + TensorShape sizes{}; + TensorShape strides{}; + size_t elems = 1; + bool changed = input.ndims() != out_rank; + const int64_t rank_offset = out_rank - input.ndims(); + + for (int64_t dim = 0; dim < out_rank; ++dim) { + const int64_t input_dim = dim - rank_offset; + const int64_t old_size = input_dim >= 0 ? input.size(input_dim) : 1; + const int64_t old_stride = + input_dim >= 0 ? input.strides()[input_dim] : 0; + + ASSERT(old_size == out_shape[dim] || old_size == 1, + fmt::format("Cannot broadcast shape {} to {}", + shape_to_string(input.dims(), input.ndims()), + shape_to_string(out_shape, out_rank))); + + sizes[dim] = out_shape[dim]; + strides[dim] = old_size == out_shape[dim] ? old_stride : 0; + elems *= static_cast(sizes[dim]); + changed |= input_dim != dim || old_size != out_shape[dim] || + old_stride != strides[dim]; + } + + if (!changed) { + return input; + } + + auto view = std::make_shared(*input.impl()); + view->sizes = sizes; + view->strides = strides; + view->ndim = out_rank; + view->elems = elems; + view->expanded = true; + return Tensor(view); +} + +TensorShape broadcast_shape(const Tensor &lhs, const Tensor &rhs, + const char *op_name) { + TensorShape out_shape{}; + const int64_t out_rank = std::max(lhs.ndims(), rhs.ndims()); + const int64_t lhs_offset = out_rank - lhs.ndims(); + const int64_t rhs_offset = out_rank - rhs.ndims(); + + for (int64_t dim = 0; dim < out_rank; ++dim) { + const int64_t lhs_dim = dim - lhs_offset; + const int64_t rhs_dim = dim - rhs_offset; + const int64_t lhs_size = lhs_dim >= 0 ? lhs.size(lhs_dim) : 1; + const int64_t rhs_size = rhs_dim >= 0 ? rhs.size(rhs_dim) : 1; + + ASSERT(lhs_size == rhs_size || lhs_size == 1 || rhs_size == 1, + fmt::format("Unable to {} non-broadcastable Tensors! {} and {}", + op_name, shape_to_string(lhs.dims(), lhs.ndims()), + shape_to_string(rhs.dims(), rhs.ndims()))); + + out_shape[dim] = std::max(lhs_size, rhs_size); + } + + return out_shape; +} + +using ContiguousBinaryLaunch = void (*)(float *, float *, float *, size_t); +using StridedBinaryLaunch = void (*)(void *, void *, void *, const StrideInfo &, + size_t); +using CpuBinaryOp = float (*)(float, float); + +float add_values(float lhs, float rhs) { return lhs + rhs; } +float sub_values(float lhs, float rhs) { return lhs - rhs; } +float mul_values(float lhs, float rhs) { return lhs * rhs; } +float div_values(float lhs, float rhs) { return lhs / rhs; } + +void compute_binary_offsets(size_t idx, const TensorShape &shape, + const TensorShape &lhs_strides, + const TensorShape &rhs_strides, int64_t rank, + int64_t &lhs_offset, int64_t &rhs_offset) { + lhs_offset = 0; + rhs_offset = 0; + + for (int64_t dim = rank - 1; dim >= 0; --dim) { + const int64_t coord = idx % shape[dim]; + idx /= shape[dim]; + lhs_offset += coord * lhs_strides[dim]; + rhs_offset += coord * rhs_strides[dim]; + } +} + +void copy_shape_to_kernel_array(const TensorShape &shape, int64_t *out, + int64_t rank) { + for (int64_t dim = 0; dim < rank; ++dim) { + out[dim] = shape[dim]; + } +} + +Tensor binary_tensor_op(const Tensor &lhs, const Tensor &rhs, + const char *op_name, + CpuBinaryOp cpu_op, + ContiguousBinaryLaunch launch_contiguous, + StridedBinaryLaunch launch_strided) { + ASSERT(lhs.device() == rhs.device(), + fmt::format("Device mismatch! {} and {}", get_device_name(lhs.device()), + get_device_name(rhs.device()))); + ASSERT(lhs.dtype() == rhs.dtype(), + fmt::format("DType mismatch! {} and {}", get_name(lhs.dtype()), + get_name(rhs.dtype()))); + + const int64_t out_rank = std::max(lhs.ndims(), rhs.ndims()); + TensorShape out_shape = broadcast_shape(lhs, rhs, op_name); + + Tensor lhs_view = make_broadcast_view(lhs, out_shape, out_rank); + Tensor rhs_view = make_broadcast_view(rhs, out_shape, out_rank); + + Tensor out = empty(out_shape.data(), out_rank, lhs.dtype(), lhs.device(), + lhs.requires_grad() || rhs.requires_grad()); + + if (lhs.device() == Device::CPU) { + const auto *lhs_data = static_cast(lhs_view.data()); + const auto *rhs_data = static_cast(rhs_view.data()); + auto *out_data = static_cast(out.data()); + + for (size_t idx = 0; idx < out.numel(); ++idx) { + int64_t lhs_offset = 0; + int64_t rhs_offset = 0; + compute_binary_offsets(idx, out_shape, lhs_view.strides(), + rhs_view.strides(), out_rank, lhs_offset, + rhs_offset); + out_data[idx] = cpu_op(lhs_data[lhs_offset], rhs_data[rhs_offset]); + } + return out; + } + + if (is_dense_contiguous(lhs_view) && is_dense_contiguous(rhs_view)) { + launch_contiguous(static_cast(out.data()), + static_cast(lhs_view.data()), + static_cast(rhs_view.data()), out.numel()); + return out; + } + + StrideInfo stride_info{}; + stride_info.rank = out_rank; + for (int64_t dim = 0; dim < out_rank; ++dim) { + stride_info.output_size[dim] = out_shape[dim]; + stride_info.a_stride[dim] = lhs_view.strides()[dim]; + stride_info.b_stride[dim] = rhs_view.strides()[dim]; + } + + launch_strided(out.data(), lhs_view.data(), rhs_view.data(), stride_info, + out.numel()); + return out; +} + +void append_tensor_values(fmt::memory_buffer &out, const float *data, + const TensorShape &sizes, + const TensorShape &strides, int64_t rank, + int64_t dim, int64_t offset) { + if (rank == 0) { + fmt::format_to(std::back_inserter(out), "{:.4f}", data[offset]); + return; + } + + fmt::format_to(std::back_inserter(out), "["); + for (int64_t i = 0; i < sizes[dim]; ++i) { + const int64_t next_offset = offset + i * strides[dim]; + if (dim == rank - 1) { + fmt::format_to(std::back_inserter(out), "{:.4f}", data[next_offset]); + } else { + append_tensor_values(out, data, sizes, strides, rank, dim + 1, + next_offset); + } + + if (i != sizes[dim] - 1) { + fmt::format_to(std::back_inserter(out), ", "); + } + } + fmt::format_to(std::back_inserter(out), "]"); +} + } // namespace template @@ -74,6 +284,10 @@ Storage::~Storage() { */ TensorImpl::TensorImpl(const int64_t *dims, int64_t rank, DataType type) { + ASSERT(rank <= static_cast(kMaxTensorDims), + fmt::format("Tensor rank {} exceeds max rank {}", rank, + kMaxTensorDims)); + for (size_t d = 0; d < rank; ++d) { sizes[d] = dims[d]; elems *= dims[d]; @@ -141,13 +355,9 @@ void *Tensor::data() const noexcept { size_t Tensor::numel() const noexcept { return impl_->elems; } -const std::array &Tensor::dims() const noexcept { - return impl_->sizes; -} +const TensorShape &Tensor::dims() const noexcept { return impl_->sizes; } -const std::array &Tensor::strides() const noexcept { - return impl_->strides; -} +const TensorShape &Tensor::strides() const noexcept { return impl_->strides; } void Tensor::print() const { if (!initialized()) { @@ -155,13 +365,14 @@ void Tensor::print() const { } else { auto &t = *impl(); fmt::print( - "Tensor: [Refcount: {} addr: {} Rank: {} dim({}, {}, {}) " - "strides({}, {}, {}) " + "Tensor: [Refcount: {} addr: {} Rank: {} dim({}) " + "strides({}) " "dtype:{} requires_grad:{}]\n\t Storage [Refcount: {} addr: {}]\n", - impl_.use_count(), fmt::ptr(impl_.get()), t.ndim, t.sizes[0], - t.sizes[1], t.sizes[2], t.strides[0], t.strides[1], t.strides[2], - get_name(t.dtype), requires_grad(), t.storage.use_count(), - t.storage->ptr); + impl_.use_count(), fmt::ptr(impl_.get()), t.ndim, + fmt::join(t.sizes.begin(), t.sizes.begin() + t.ndim, ", "), + fmt::join(t.strides.begin(), t.strides.begin() + t.ndim, ", "), + get_name(t.dtype), + requires_grad(), t.storage.use_count(), t.storage->ptr); } } @@ -173,9 +384,6 @@ std::string Tensor::to_string() const { return "[]"; } - ASSERT(ndims() <= 3, - fmt::format("Tensor::print_elms unsupported ndims=={}", ndims())); - // Could be expensive auto t = cpu(); const float *raw_data = static_cast(t.data()); @@ -185,48 +393,9 @@ std::string Tensor::to_string() const { fmt::memory_buffer out; - if (ndims() == 1) { - fmt::format_to(std::back_inserter(out), "Tensor: (["); - for (int64_t i = 0; i < sizes[0]; ++i) { - fmt::format_to(std::back_inserter(out), "{:.4f}{}", - raw_data[i * stride[0]], i == sizes[0] - 1 ? "" : ", "); - } - fmt::format_to(std::back_inserter(out), "])\n"); - } else if (ndims() == 2) { - fmt::format_to(std::back_inserter(out), "Tensor: (["); - for (int64_t i = 0; i < sizes[0]; ++i) { - fmt::format_to(std::back_inserter(out), "["); - for (int64_t j = 0; j < sizes[1]; ++j) { - fmt::format_to(std::back_inserter(out), "{:.4f}{}", - raw_data[i * stride[0] + j * stride[1]], - j == sizes[1] - 1 ? "" : ", "); - } - fmt::format_to(std::back_inserter(out), "{}", - i == sizes[0] - 1 ? "]" : "],\n "); - } - fmt::format_to(std::back_inserter(out), "])\n"); - } else if (ndims() == 3) { - fmt::format_to(std::back_inserter(out), "Tensor: (["); - for (int64_t i = 0; i < sizes[0]; ++i) { - fmt::format_to(std::back_inserter(out), "["); - for (int64_t j = 0; j < sizes[1]; ++j) { - fmt::format_to(std::back_inserter(out), "["); - for (int64_t k = 0; k < sizes[2]; ++k) { - fmt::format_to( - std::back_inserter(out), "{:.4f}{}", - raw_data[k * stride[2] + j * stride[1] + i * stride[0]], - k == sizes[2] - 1 ? "" : ", "); - } - fmt::format_to(std::back_inserter(out), "{}", - j == sizes[1] - 1 ? "]" : "],\n "); - } - fmt::format_to(std::back_inserter(out), "{}", - i == sizes[0] - 1 ? "]" : "],\n\n "); - } - fmt::format_to(std::back_inserter(out), "])\n"); - } - - fmt::format_to(std::back_inserter(out), "])\n"); + fmt::format_to(std::back_inserter(out), "Tensor: ("); + append_tensor_values(out, raw_data, sizes, stride, ndims(), 0, 0); + fmt::format_to(std::back_inserter(out), ")\n"); return fmt::to_string(out); } @@ -244,53 +413,9 @@ Tensor Tensor::sum(int64_t dim, bool keep_dim) const { Tensor Tensor::add(float scalar) const { return add(full_like(*this, scalar)); } Tensor Tensor::mul(const Tensor &other) const { - std::array out_sz = {0, 0, 0}; - bool expand_me = false; - bool expand_other = false; - - int64_t out_rank = 0; - for (int i = 0; i < 3; ++i) { - const auto my_size = size(i); - const auto other_size = other.size(i); - ASSERT( - my_size == other_size or (my_size == 1 or my_size == 0) or - (other_size == 1 or other_size == 0), - fmt::format("Unable to multiply non-broadcastable Tensors! [{},{},{}] " - "and [{},{},{}]", - size(0), size(1), size(2), other.size(0), other.size(1), - other.size(2))); - - out_sz[i] = std::max(impl()->sizes[i], other.impl()->sizes[i]); - - if (out_sz[i] > 0) { - out_rank++; - } - - expand_me |= out_sz[i] != my_size; - expand_other |= out_sz[i] != other_size; - } - - auto me_alias = expand_me ? expand(out_sz) : *this; - auto other_alias = expand_other ? other.expand(out_sz) : other; - - Tensor out = empty(out_sz.data(), out_rank, dtype(), device(), - requires_grad() || other.requires_grad()); - - if (!expand_me and !expand_other) { - launch_mul(static_cast(out.data()), static_cast(data()), - static_cast(other.data()), out.numel()); - } else { - StrideInfo s{}; - s.rank = out_rank; - for (int i = 0; i < s.rank; ++i) { - s.output_size[i] = out_sz[i]; - s.a_stride[i] = me_alias.impl()->strides[i]; - s.b_stride[i] = other_alias.impl()->strides[i]; - } - - launch_mul_strided(out.data(), me_alias.data(), other_alias.data(), s, - out.numel()); - } + Tensor out = + binary_tensor_op(*this, other, "multiply", mul_values, launch_mul, + launch_mul_strided); SetupAutograd(*this, other, out); return out; @@ -303,105 +428,18 @@ Tensor Tensor::matmul(const Tensor &other) const { } Tensor Tensor::add(const Tensor &other) const { - - std::array out_sz = {0, 0, 0}; - bool expand_me = false; - bool expand_other = false; - - int64_t out_rank = 0; - for (int i = 0; i < 3; ++i) { - const auto my_size = size(i); - const auto other_size = other.size(i); - ASSERT(my_size == other_size or (my_size == 1 or my_size == 0) or - (other_size == 1 or other_size == 0), - fmt::format("Unable to add non-broadcastable Tensors! [{},{},{}] " - "and [{},{},{}]", - size(0), size(1), size(2), other.size(0), other.size(1), - other.size(2))); - - out_sz[i] = std::max(impl()->sizes[i], other.impl()->sizes[i]); - - if (out_sz[i] > 0) { - out_rank++; - } - - expand_me |= out_sz[i] != my_size; - expand_other |= out_sz[i] != other_size; - } - - auto me_alias = expand_me ? expand(out_sz) : *this; - auto other_alias = expand_other ? other.expand(out_sz) : other; - - Tensor out = empty(out_sz.data(), out_rank, dtype(), device(), - requires_grad() || other.requires_grad()); - - if (!expand_me and !expand_other) { - launch_add(static_cast(out.data()), static_cast(data()), - static_cast(other.data()), out.numel()); - } else { - StrideInfo s{}; - s.rank = out_rank; - for (int i = 0; i < s.rank; ++i) { - s.output_size[i] = out_sz[i]; - s.a_stride[i] = me_alias.impl()->strides[i]; - s.b_stride[i] = other_alias.impl()->strides[i]; - } - - launch_add_strided(out.data(), me_alias.data(), other_alias.data(), s, - out.numel()); - } + Tensor out = + binary_tensor_op(*this, other, "add", add_values, launch_add, + launch_add_strided); SetupAutograd(*this, other, out); return out; } Tensor Tensor::sub(const Tensor &other) const { - std::array out_sz = {0, 0, 0}; - bool expand_me = false; - bool expand_other = false; - - int64_t out_rank = 0; - for (int i = 0; i < 3; ++i) { - const auto my_size = size(i); - const auto other_size = other.size(i); - ASSERT(my_size == other_size or (my_size == 1 or my_size == 0) or - (other_size == 1 or other_size == 0), - fmt::format("Unable to add non-broadcastable Tensors! [{},{},{}] " - "and [{},{},{}]", - size(0), size(1), size(2), other.size(0), other.size(1), - other.size(2))); - - out_sz[i] = std::max(impl()->sizes[i], other.impl()->sizes[i]); - - if (out_sz[i] > 0) { - out_rank++; - } - - expand_me |= out_sz[i] != my_size; - expand_other |= out_sz[i] != other_size; - } - - auto me_alias = expand_me ? expand(out_sz) : *this; - auto other_alias = expand_other ? other.expand(out_sz) : other; - - Tensor out = empty(out_sz.data(), out_rank, dtype(), device(), - requires_grad() || other.requires_grad()); - - if (!expand_me and !expand_other) { - launch_sub(static_cast(out.data()), static_cast(data()), - static_cast(other.data()), out.numel()); - } else { - StrideInfo s{}; - s.rank = out_rank; - for (int i = 0; i < s.rank; ++i) { - s.output_size[i] = out_sz[i]; - s.a_stride[i] = me_alias.impl()->strides[i]; - s.b_stride[i] = other_alias.impl()->strides[i]; - } - - launch_sub_strided(out.data(), me_alias.data(), other_alias.data(), s, - out.numel()); - } + Tensor out = + binary_tensor_op(*this, other, "subtract", sub_values, launch_sub, + launch_sub_strided); SetupAutograd(*this, other, out); return out; @@ -414,52 +452,9 @@ Tensor Tensor::rsub(float scalar) const { } Tensor Tensor::div(const Tensor &other) const { - std::array out_sz = {0, 0, 0}; - bool expand_me = false; - bool expand_other = false; - - int64_t out_rank = 0; - for (int i = 0; i < 3; ++i) { - const auto my_size = size(i); - const auto other_size = other.size(i); - ASSERT(my_size == other_size or (my_size == 1 or my_size == 0) or - (other_size == 1 or other_size == 0), - fmt::format("Unable to divide non-broadcastable Tensors! [{},{},{}] " - "and [{},{},{}]", - size(0), size(1), size(2), other.size(0), other.size(1), - other.size(2))); - - out_sz[i] = std::max(impl()->sizes[i], other.impl()->sizes[i]); - - if (out_sz[i] > 0) { - out_rank++; - } - - expand_me |= out_sz[i] != my_size; - expand_other |= out_sz[i] != other_size; - } - - auto me_alias = expand_me ? expand(out_sz) : *this; - auto other_alias = expand_other ? other.expand(out_sz) : other; - - Tensor out = empty(out_sz.data(), out_rank, dtype(), device(), - requires_grad() || other.requires_grad()); - - if (!expand_me and !expand_other) { - launch_div(static_cast(out.data()), static_cast(data()), - static_cast(other.data()), out.numel()); - } else { - StrideInfo s{}; - s.rank = out_rank; - for (int i = 0; i < s.rank; ++i) { - s.output_size[i] = out_sz[i]; - s.a_stride[i] = me_alias.impl()->strides[i]; - s.b_stride[i] = other_alias.impl()->strides[i]; - } - - launch_div_strided(out.data(), me_alias.data(), other_alias.data(), s, - out.numel()); - } + Tensor out = + binary_tensor_op(*this, other, "divide", div_values, launch_div, + launch_div_strided); SetupAutograd(*this, other, out); return out; @@ -491,45 +486,11 @@ Tensor Tensor::transpose(int d0, int d1) const { return return_tensor; } -Tensor Tensor::expand(const std::array &new_sz) const { - const auto &old = impl()->sizes; - const int64_t rank = impl()->ndim; - - // check broadcast-compatibility and build new strides - std::array ns = old; - std::array st = impl()->strides; - - size_t elems = 1; - for (int i = 0; i < rank; ++i) { - if (old[i] == new_sz[i]) { - ns[i] = new_sz[i]; - } else { - ASSERT(old[i] == 1, - fmt::format("expand: non-broadcastable dim {}", old[i])); - ns[i] = new_sz[i]; - st[i] = 0; - } - - elems *= ns[i]; - } - - // make view - auto v = std::make_shared(); - v->sizes = ns; - v->strides = st; - v->dtype = dtype(); - v->ndim = rank; - v->elems = elems; - v->storage = impl()->storage; - v->requires_grad = requires_grad(); - v->expanded = true; - - // share autograd meta - if (impl()->grad) { - v->grad = impl()->grad; - } - - return Tensor(v); +Tensor Tensor::expand(const TensorShape &new_sz) const { + const int64_t new_rank = infer_rank(new_sz); + ASSERT(new_rank > 0 || ndims() == 0, + "expand requires at least one non-zero dimension"); + return make_broadcast_view(*this, new_sz, new_rank); } Tensor Tensor::cuda() const { @@ -725,22 +686,26 @@ Tensor sum(const Tensor &t, int64_t dim, bool keep_dim) { dims[2] = std::max(dims[2], 1l); StrideAndSize s_input; - s_input.size = dims; - s_input.stride = t.strides(); + for (int64_t i = 0; i < t.ndims(); ++i) { + s_input.size[i] = dims[i]; + s_input.stride[i] = t.strides()[i]; + } s_input.rank = t.ndims(); StrideAndSize s_output; - s_output.size = new_tensor.dims(); - s_output.stride = new_tensor.strides(); + for (int64_t i = 0; i < new_tensor.ndims(); ++i) { + s_output.size[i] = new_tensor.dims()[i]; + s_output.stride[i] = new_tensor.strides()[i]; + } s_output.rank = new_tensor.ndims(); if (dim == 0) { - launch_sum_dim0(dst, srcp, s_input, s_output); + launch_sum_dim(dst, srcp, s_input, s_output, 0); } else if (dim == 1) { - launch_sum_dim1(dst, srcp, s_input, s_output); + launch_sum_dim(dst, srcp, s_input, s_output, 1); } else { // dim==2 - launch_sum_dim2(dst, srcp, s_input, s_output); + launch_sum_dim(dst, srcp, s_input, s_output, 2); } return new_tensor; @@ -835,6 +800,10 @@ Tensor &operator/=(Tensor &l, float scalar) { Tensor empty(const int64_t *dims, size_t rank, DataType t, Device d, bool requires_grad) { + ASSERT(rank <= kMaxTensorDims, + fmt::format("Tensor rank {} exceeds max rank {}", rank, + kMaxTensorDims)); + auto storage = std::make_shared(); float *ptr; From 254f4f5af11adf9dd1bcf710117ee08c76193710 Mon Sep 17 00:00:00 2001 From: Jacob Domagala Date: Wed, 12 Nov 2025 12:03:00 +0100 Subject: [PATCH 5/9] [#80]: Add higher-rank sum fallback --- src/sum.cu | 98 ++++++++++++++++++++++++++++++++++++++++++++++++-- src/tensor.cpp | 40 +++++++++------------ 2 files changed, 112 insertions(+), 26 deletions(-) diff --git a/src/sum.cu b/src/sum.cu index 3bf2914..78790df 100644 --- a/src/sum.cu +++ b/src/sum.cu @@ -3,6 +3,25 @@ namespace smollnet { +namespace { + +size_t shape_numel(const int64_t *shape, int64_t rank) { + size_t total = 1; + for (int64_t dim = 0; dim < rank; ++dim) { + total *= static_cast(shape[dim]); + } + return total; +} + +StrideAndSize padded_rank3(StrideAndSize shape) { + for (int64_t dim = shape.rank; dim < 3; ++dim) { + shape.size[dim] = 1; + } + return shape; +} + +} // namespace + template __global__ void warp_level_sum(const float *__restrict__ in, float *__restrict__ out, @@ -25,7 +44,7 @@ warp_level_sum(const float *__restrict__ in, float *__restrict__ out, idx = depth * s0_in + row * s1_in + col; out_idx = depth * s0_out + row * s1_out; in_bounds = (idx < n and col < dim_size); - } else if (MAJOR == COL_MAJOR) { + } else if constexpr (MAJOR == COL_MAJOR) { col = blockIdx.y; row = blockIdx.x * blockDim.x + threadIdx.x; @@ -58,7 +77,7 @@ warp_level_sum(const float *__restrict__ in, float *__restrict__ out, if (threadIdx.x == 0) { float acc = 0.0f; - #pragma unroll +#pragma unroll for (int i = 0; i < BLOCK_DIM / 32; ++i) { acc += sMem[i]; } @@ -313,4 +332,79 @@ void launch_sum_dim2(void *out, void *in, const StrideAndSize &s_input, CHECK_CUDA(cudaGetLastError()); } +__global__ void generic_sum_dim_kernel(const float *__restrict__ in, + float *__restrict__ out, + const StrideAndSize s_input, + const StrideAndSize s_output, + const int64_t reduce_dim, + const size_t total) { + size_t linear = blockIdx.x * blockDim.x + threadIdx.x; + if (linear >= total) { + return; + } + + size_t remaining = linear; + int64_t input_offset = 0; + int64_t output_offset = 0; + + for (int64_t dim = s_input.rank - 1; dim >= 0; --dim) { + const int64_t coord = remaining % s_input.size[dim]; + remaining /= s_input.size[dim]; + + input_offset += coord * s_input.stride[dim]; + + if (dim == reduce_dim) { + continue; + } + + int64_t output_dim = dim; + if (s_output.rank != s_input.rank && dim > reduce_dim) { + output_dim = dim - 1; + } + + output_offset += coord * s_output.stride[output_dim]; + } + + atomicAdd(out + output_offset, in[input_offset]); +} + +void launch_generic_sum_dim(void *out, void *in, const StrideAndSize &s_input, + const StrideAndSize &s_output, int64_t dim) { + const size_t total = shape_numel(s_input.size, s_input.rank); + if (total == 0) { + return; + } + + constexpr int block = 256; + const int grid = static_cast((total + block - 1) / block); + + generic_sum_dim_kernel<<>>( + static_cast(in), static_cast(out), s_input, + s_output, dim, total); + + CHECK_CUDA(cudaGetLastError()); +} + +void launch_sum_dim(void *out, void *in, const StrideAndSize &s_input, + const StrideAndSize &s_output, int64_t dim) { + if (s_input.rank <= 3) { + const StrideAndSize opt_input = padded_rank3(s_input); + const StrideAndSize opt_output = padded_rank3(s_output); + + if (dim == 0) { + launch_sum_dim0(out, in, opt_input, opt_output); + return; + } + if (dim == 1) { + launch_sum_dim1(out, in, opt_input, opt_output); + return; + } + + launch_sum_dim2(out, in, opt_input, opt_output); + return; + } + + launch_generic_sum_dim(out, in, s_input, s_output, dim); +} + } // namespace smollnet diff --git a/src/tensor.cpp b/src/tensor.cpp index b9be32a..c065919 100644 --- a/src/tensor.cpp +++ b/src/tensor.cpp @@ -666,8 +666,13 @@ Tensor sum(const Tensor &t, int64_t dim, bool keep_dim) { "Tensor sum(tensor,dim,keep_dim): invalid dim={} t.ndims()={}", dim, t.ndims())); + ASSERT(dim >= 0, + fmt::format( + "Tensor sum(tensor,dim,keep_dim): invalid dim={} t.ndims()={}", + dim, t.ndims())); + // build output shape - int64_t out_dims[3] = {0, 0, 0}; + TensorShape out_dims{}; for (int64_t i = 0, o = 0; i < t.ndims(); ++i) { if (i != dim) { out_dims[o++] = dims[i]; @@ -677,36 +682,23 @@ Tensor sum(const Tensor &t, int64_t dim, bool keep_dim) { } Tensor new_tensor = - zeros(out_dims, new_rank, data_type, device, t.requires_grad()); + zeros(out_dims.data(), new_rank, data_type, device, t.requires_grad()); auto *srcp = t.data(); auto *dst = new_tensor.data(); - dims[0] = std::max(dims[0], 1l); - dims[1] = std::max(dims[1], 1l); - dims[2] = std::max(dims[2], 1l); - - StrideAndSize s_input; - for (int64_t i = 0; i < t.ndims(); ++i) { - s_input.size[i] = dims[i]; - s_input.stride[i] = t.strides()[i]; - } + + StrideAndSize s_input{}; s_input.rank = t.ndims(); + copy_shape_to_kernel_array(dims, s_input.size, s_input.rank); + copy_shape_to_kernel_array(t.strides(), s_input.stride, s_input.rank); - StrideAndSize s_output; - for (int64_t i = 0; i < new_tensor.ndims(); ++i) { - s_output.size[i] = new_tensor.dims()[i]; - s_output.stride[i] = new_tensor.strides()[i]; - } + StrideAndSize s_output{}; s_output.rank = new_tensor.ndims(); + copy_shape_to_kernel_array(new_tensor.dims(), s_output.size, s_output.rank); + copy_shape_to_kernel_array(new_tensor.strides(), s_output.stride, + s_output.rank); - if (dim == 0) { - launch_sum_dim(dst, srcp, s_input, s_output, 0); - } else if (dim == 1) { - launch_sum_dim(dst, srcp, s_input, s_output, 1); - } else { - // dim==2 - launch_sum_dim(dst, srcp, s_input, s_output, 2); - } + launch_sum_dim(dst, srcp, s_input, s_output, dim); return new_tensor; } From 3228ca142847016369b5904e4609898221b4fcf3 Mon Sep 17 00:00:00 2001 From: Jacob Domagala Date: Wed, 12 Nov 2025 12:04:00 +0100 Subject: [PATCH 6/9] [#80]: Validate SGD gradients by rank --- src/sgd.cpp | 19 +++++++++---------- 1 file changed, 9 insertions(+), 10 deletions(-) diff --git a/src/sgd.cpp b/src/sgd.cpp index a27a734..7c99b8a 100644 --- a/src/sgd.cpp +++ b/src/sgd.cpp @@ -8,17 +8,16 @@ namespace smollnet { void SGD::step() const { for (const auto &p : params_) { - ASSERT( - p.size(0) == p.grad().size(0), - fmt::format("Size 0 mismatch!: {} vs {}", p.size(0), p.grad().size(0))); - ASSERT( - p.size(1) == p.grad().size(1), - fmt::format("Size 1 mismatch!: {} vs {}", p.size(1), p.grad().size(1))); - ASSERT( - p.size(2) == p.grad().size(2), - fmt::format("Size 2 mismatch!: {} vs {}", p.size(2), p.grad().size(2))); + Tensor grad = p.grad(); + ASSERT(p.ndims() == grad.ndims(), + fmt::format("Rank mismatch!: {} vs {}", p.ndims(), grad.ndims())); + for (int64_t dim = 0; dim < p.ndims(); ++dim) { + ASSERT(p.size(dim) == grad.size(dim), + fmt::format("Size {} mismatch!: {} vs {}", dim, p.size(dim), + grad.size(dim))); + } - launch_sgd_update(p.data(), p.grad().data(), lr_, p.numel()); + launch_sgd_update(p.data(), grad.data(), lr_, p.numel()); } } From e04fdb78e8dfd351eef42c4293298c7d1e4e8fb6 Mon Sep 17 00:00:00 2001 From: Jacob Domagala Date: Wed, 12 Nov 2025 12:05:00 +0100 Subject: [PATCH 7/9] [#80]: Make matmul explicitly 2D --- src/tensor.cpp | 34 ++++++++++++---------------------- 1 file changed, 12 insertions(+), 22 deletions(-) diff --git a/src/tensor.cpp b/src/tensor.cpp index c065919..b452548 100644 --- a/src/tensor.cpp +++ b/src/tensor.cpp @@ -558,20 +558,16 @@ Tensor matmul(const Tensor &l, const Tensor &r) { l.ndims(), r.ndims())); // TODO: allow for broadcast - ASSERT(l.dims().size() == r.dims().size(), - fmt::format("{} vs {}", l.dims().size(), r.dims().size())); - - if (l.ndims() == 2) { - ASSERT(l.dims()[1] == r.dims()[0], - fmt::format("Incorrect matrix size! lhs number of rows ({}) not " - "equal to rhs number of cols ({})", - l.dims()[1], r.dims()[0])); - } else { - ASSERT(l.dims()[2] == r.dims()[1], - fmt::format("Incorrect matrix size! lhs number of rows ({}) not " - "equal to rhs number of cols ({})", - l.dims()[2], r.dims()[1])); - } + ASSERT(l.ndims() == r.ndims(), + fmt::format("Matmul rank mismatch: {} vs {}", l.ndims(), r.ndims())); + ASSERT(l.ndims() == 2, + fmt::format("Matmul currently supports 2D tensors, got rank {}", + l.ndims())); + + ASSERT(l.dims()[1] == r.dims()[0], + fmt::format("Incorrect matrix size! lhs number of rows ({}) not " + "equal to rhs number of cols ({})", + l.dims()[1], r.dims()[0])); ASSERT(l.device() == r.device(), fmt::format("Device mismatch! {} and {}", get_device_name(l.device()), @@ -581,32 +577,26 @@ Tensor matmul(const Tensor &l, const Tensor &r) { Tensor new_tensor = empty({l.dims()[0], r.dims()[1]}, l.dtype(), l.device(), needs_grad); - StrideInfo stride_info; + StrideInfo stride_info{}; stride_info.output_size[0] = new_tensor.size(0); stride_info.output_size[1] = new_tensor.size(1); const auto &l_strides = l.strides(); - ; stride_info.a_stride[0] = l_strides[0]; stride_info.a_stride[1] = l_strides[1]; - stride_info.a_stride[2] = l_strides[2]; const auto &r_strides = r.strides(); - ; stride_info.b_stride[0] = r_strides[0]; stride_info.b_stride[1] = r_strides[1]; - stride_info.b_stride[2] = r_strides[2]; stride_info.rank = new_tensor.ndims(); - SizeInfo size_info; + SizeInfo size_info{}; size_info.a_size[0] = l.size(0); size_info.a_size[1] = l.size(1); - size_info.a_size[2] = l.size(2); size_info.b_size[0] = r.size(0); size_info.b_size[1] = r.size(1); - size_info.b_size[2] = r.size(2); launch_matmul(new_tensor.data(), l.data(), r.data(), stride_info, size_info, new_tensor.numel()); From feb735023e820cd37bb6cfbaf409fe82d2940357 Mon Sep 17 00:00:00 2001 From: Jacob Domagala Date: Wed, 12 Nov 2025 12:06:00 +0100 Subject: [PATCH 8/9] [#80]: Allow Python tensor creation from shape lists --- python/bindings.cpp | 47 ++++++++++++++++++++++++++++++++++++++++++++- 1 file changed, 46 insertions(+), 1 deletion(-) diff --git a/python/bindings.cpp b/python/bindings.cpp index a89a4db..7914f8e 100644 --- a/python/bindings.cpp +++ b/python/bindings.cpp @@ -4,6 +4,9 @@ #include #include +#include +#include + using namespace pybind11::literals; namespace py = pybind11; @@ -11,6 +14,20 @@ namespace py = pybind11; template void bind_tensor_creation_overloads(pybind11::module &m, const char *func_name, FuncType &&func) { + m.def( + func_name, + [func](const std::vector &dims, + smollnet::DataType dtype = smollnet::DataType::f32, + smollnet::Device device = smollnet::Device::CUDA, + bool requires_grad = false) { + if (dims.size() > smollnet::kMaxTensorDims) { + throw std::invalid_argument("Tensor rank exceeds kMaxTensorDims"); + } + return func(dims.data(), dims.size(), dtype, device, requires_grad); + }, + "dims"_a, "dtype"_a = smollnet::DataType::f32, + "device"_a = smollnet::Device::CUDA, "requires_grad"_a = false); + // 1D version m.def( func_name, @@ -55,6 +72,11 @@ void bind_tensor_creation_overloads(pybind11::module &m, const char *func_name, // Function objects for each tensor creation function struct RandFunctor { + auto operator()(const int64_t *dims, size_t rank, smollnet::DataType dtype, + smollnet::Device device, bool requires_grad) const { + return smollnet::rand(dims, rank, dtype, device, requires_grad); + } + template auto operator()(const int64_t (&dims)[N], smollnet::DataType dtype, smollnet::Device device, bool requires_grad) const { @@ -63,6 +85,11 @@ struct RandFunctor { }; struct ZerosFunctor { + auto operator()(const int64_t *dims, size_t rank, smollnet::DataType dtype, + smollnet::Device device, bool requires_grad) const { + return smollnet::zeros(dims, rank, dtype, device, requires_grad); + } + template auto operator()(const int64_t (&dims)[N], smollnet::DataType dtype, smollnet::Device device, bool requires_grad) const { @@ -71,6 +98,11 @@ struct ZerosFunctor { }; struct OnesFunctor { + auto operator()(const int64_t *dims, size_t rank, smollnet::DataType dtype, + smollnet::Device device, bool requires_grad) const { + return smollnet::ones(dims, rank, dtype, device, requires_grad); + } + template auto operator()(const int64_t (&dims)[N], smollnet::DataType dtype, smollnet::Device device, bool requires_grad) const { @@ -79,6 +111,11 @@ struct OnesFunctor { }; struct EmptyFunctor { + auto operator()(const int64_t *dims, size_t rank, smollnet::DataType dtype, + smollnet::Device device, bool requires_grad) const { + return smollnet::empty(dims, rank, dtype, device, requires_grad); + } + template auto operator()(const int64_t (&dims)[N], smollnet::DataType dtype, smollnet::Device device, bool requires_grad) const { @@ -155,7 +192,15 @@ PYBIND11_MODULE(smollnet, m) { .def("matmul", &smollnet::Tensor::matmul) .def("transpose", &smollnet::Tensor::transpose) - .def("expand", &smollnet::Tensor::expand) + .def("expand", + [](const smollnet::Tensor &tensor, const std::vector &dims) { + if (dims.size() > smollnet::kMaxTensorDims) { + throw std::invalid_argument("Tensor rank exceeds kMaxTensorDims"); + } + smollnet::TensorShape shape{}; + std::copy(dims.begin(), dims.end(), shape.begin()); + return tensor.expand(shape); + }) .def("cuda", &smollnet::Tensor::cuda) .def("cpu", &smollnet::Tensor::cpu) From 4e4571e11a3c29915f128b58607b2fa49c1758cf Mon Sep 17 00:00:00 2001 From: Jacob Domagala Date: Wed, 12 Nov 2025 12:07:00 +0100 Subject: [PATCH 9/9] [#80]: Fix tensor narrowing diagnostics --- src/tensor.cpp | 7 ++++--- src/tensor.hpp | 2 +- 2 files changed, 5 insertions(+), 4 deletions(-) diff --git a/src/tensor.cpp b/src/tensor.cpp index b452548..b4c34a2 100644 --- a/src/tensor.cpp +++ b/src/tensor.cpp @@ -135,9 +135,10 @@ void compute_binary_offsets(size_t idx, const TensorShape &shape, lhs_offset = 0; rhs_offset = 0; + int64_t remaining = static_cast(idx); for (int64_t dim = rank - 1; dim >= 0; --dim) { - const int64_t coord = idx % shape[dim]; - idx /= shape[dim]; + const int64_t coord = remaining % shape[dim]; + remaining /= shape[dim]; lhs_offset += coord * lhs_strides[dim]; rhs_offset += coord * rhs_strides[dim]; } @@ -341,7 +342,7 @@ Tensor Tensor::grad() const noexcept { AutogradMeta *Tensor::autograd() const noexcept { return impl()->grad.get(); } -int64_t Tensor::size(int d) const noexcept { return impl()->sizes[d]; } +int64_t Tensor::size(int64_t d) const noexcept { return impl()->sizes[d]; } int64_t Tensor::ndims() const noexcept { return impl()->ndim; } diff --git a/src/tensor.hpp b/src/tensor.hpp index ab86bdd..e7fe6dc 100644 --- a/src/tensor.hpp +++ b/src/tensor.hpp @@ -70,7 +70,7 @@ class Tensor { bool requires_grad() const noexcept; Tensor grad() const noexcept; AutogradMeta *autograd() const noexcept; - int64_t size(int d) const noexcept; + int64_t size(int64_t d) const noexcept; int64_t ndims() const noexcept; Device device() const noexcept; DataType dtype() const noexcept;