#include "tensor.h"

#include "matrix.h"

#include <float.h>

#include <math.h>

#include <stdint.h>

int buffers_shapes_match(Py_buffer const *first, Py_buffer const *second) {

if (first->ndim != second->ndim) {

PyErr_SetString(PyExc_ValueError, "Input tensor ranks don't match");

return 0;

}

for (int dimension = 0; dimension < first->ndim; ++dimension) {

if (first->shape[dimension] != second->shape[dimension]) {

PyErr_Format(

PyExc_ValueError,

"Input tensor shapes don't match at dimension %d (%zd vs %zd). " "NumKong does not support " "implicit " "shape " "broadcasting.",

dimension, first->shape[dimension], second->shape[dimension]);

return 0;

}

}

return 1;

}

size_t shared_contiguous_tail_dimensions(Py_buffer const *buffers[], size_t num_buffers, size_t num_dims) {

size_t num_contiguous_dims = 0;

for (size_t dimension = num_dims; dimension-- > 0;) {

// Compute the expected stride for this dimension if it were packed

int all_packed = 1;

for (size_t buffer_idx = 0; buffer_idx < num_buffers; ++buffer_idx) {

Py_ssize_t expected_stride = buffers[buffer_idx]->itemsize;

for (size_t inner_dim = num_dims - 1; inner_dim > dimension; --inner_dim)

expected_stride *= buffers[buffer_idx]->shape[inner_dim];

if (buffers[buffer_idx]->strides[dimension] != expected_stride) {

all_packed = 0;

break;

}

}

if (!all_packed) break;

++num_contiguous_dims;

}

return num_contiguous_dims;

}

void each_sum_recursive( //

size_t remaining_dims, size_t contiguous_tail_dims) {

// Base case: all remaining dimensions are contiguous — one kernel call

if (remaining_dims <= contiguous_tail_dims) {

size_t contiguous_elements = 1;

for (size_t dimension = 0; dimension < remaining_dims; ++dimension)

contiguous_elements *= (size_t)shape[dimension];

kernel(a_data, b_data, contiguous_elements, result_data);

return;

}

// Iterate over the outermost non-contiguous dimension, then recurse

size_t const dim_extent = (size_t)shape[0];

for (size_t position = 0; position < dim_extent; ++position) {

Py_ssize_t const signed_position = (Py_ssize_t)position;

each_sum_recursive(kernel, //

a_data + signed_position * a_strides[0], //

b_data + signed_position * b_strides[0], //

result_data + signed_position * result_strides[0], //

shape + 1, a_strides + 1, //

b_strides + 1, result_strides + 1, //

remaining_dims - 1, contiguous_tail_dims);

}

}

void each_scale_recursive( //

size_t remaining_dims, size_t contiguous_tail_dims) {

if (remaining_dims <= contiguous_tail_dims) {

size_t contiguous_elements = 1;

for (size_t dimension = 0; dimension < remaining_dims; ++dimension)

contiguous_elements *= (size_t)shape[dimension];

kernel(a_data, contiguous_elements, alpha, beta, result_data);

return;

}

size_t const dim_extent = (size_t)shape[0];

for (size_t position = 0; position < dim_extent; ++position) {

Py_ssize_t const signed_position = (Py_ssize_t)position;

each_scale_recursive(kernel, //

a_data + signed_position * a_strides[0], //

result_data + signed_position * result_strides[0], //

alpha, beta, //

shape + 1, a_strides + 1, //

result_strides + 1, //

remaining_dims - 1, contiguous_tail_dims);

}

}

void each_fma_recursive( //

size_t remaining_dims, size_t contiguous_tail_dims) {

if (remaining_dims <= contiguous_tail_dims) {

size_t contiguous_elements = 1;

for (size_t dimension = 0; dimension < remaining_dims; ++dimension)

contiguous_elements *= (size_t)shape[dimension];

kernel(a_data, b_data, c_data, contiguous_elements, alpha, beta, result_data);

return;

}

size_t const dim_extent = (size_t)shape[0];

for (size_t position = 0; position < dim_extent; ++position) {

Py_ssize_t const signed_position = (Py_ssize_t)position;

each_fma_recursive(kernel, //

a_data + signed_position * a_strides[0], //

b_data + signed_position * b_strides[0], //

c_data + signed_position * c_strides[0], //

result_data + signed_position * result_strides[0], //

alpha, beta, //

shape + 1, a_strides + 1, //

b_strides + 1, c_strides + 1, //

result_strides + 1, //

remaining_dims - 1, contiguous_tail_dims);

}

}

void each_blend_recursive( //

size_t remaining_dims, size_t contiguous_tail_dims) {

if (remaining_dims <= contiguous_tail_dims) {

size_t contiguous_elements = 1;

for (size_t dimension = 0; dimension < remaining_dims; ++dimension)

contiguous_elements *= (size_t)shape[dimension];

kernel(a_data, b_data, contiguous_elements, alpha, beta, result_data);

return;

}

size_t const dim_extent = (size_t)shape[0];

for (size_t position = 0; position < dim_extent; ++position) {

Py_ssize_t const signed_position = (Py_ssize_t)position;

each_blend_recursive(kernel, //

a_data + signed_position * a_strides[0], //

b_data + signed_position * b_strides[0], //

result_data + signed_position * result_strides[0], //

alpha, beta, //

shape + 1, a_strides + 1, //

b_strides + 1, result_strides + 1, //

remaining_dims - 1, contiguous_tail_dims);

}

}

static int tensor_is_c_contig(Tensor *tensor, size_t item_size);

static int tensor_is_f_contig(Tensor *tensor, size_t item_size);

static PyObject *tensor_read_scalar(Tensor *tensor, size_t byte_offset) {

size_t elem_size = bytes_per_dtype(tensor->dtype);

if (!elem_size) {

PyErr_SetString(PyExc_TypeError, "unsupported dtype for indexing");

return NULL;

}

nk_scalar_buffer_t buf;

memset(&buf, 0, sizeof(buf));

memcpy(&buf, tensor->data + byte_offset, elem_size);

return scalar_to_py_number(&buf, tensor->dtype);

}

static int tensor_is_c_contig(Tensor *tensor, size_t item_size) {

if (tensor->rank == 0) return 1;

Py_ssize_t expected = (Py_ssize_t)item_size;

for (size_t i = tensor->rank; i > 0; i--) {

if (tensor->strides[i - 1] != expected) return 0;

expected *= tensor->shape[i - 1];

}

return 1;

}

static int tensor_is_f_contig(Tensor *tensor, size_t item_size) {

if (tensor->rank == 0) return 1;

Py_ssize_t expected = (Py_ssize_t)item_size;

for (size_t i = 0; i < tensor->rank; i++) {

if (tensor->strides[i] != expected) return 0;

expected *= tensor->shape[i];

}

return 1;

}

static void Tensor_dealloc(PyObject *self) {

Tensor *tensor = (Tensor *)self;

Py_XDECREF(tensor->parent);

Py_TYPE(self)->tp_free(self);

}

Tensor *Tensor_new(nk_dtype_t dtype, size_t rank, Py_ssize_t const *shape) {

if (rank > NK_TENSOR_MAX_RANK) {

PyErr_Format(PyExc_ValueError, "Tensor rank %zu exceeds maximum %d", rank, NK_TENSOR_MAX_RANK);

return NULL;

}

size_t const item_size = bytes_per_dtype(dtype);

size_t total_items = 1;

for (size_t i = 0; i < rank; i++) {

if (shape[i] > 0 && total_items > SIZE_MAX / (size_t)shape[i]) {

PyErr_SetString(PyExc_OverflowError, "Tensor shape too large");

return NULL;

}

total_items *= (size_t)shape[i];

}

if (item_size > 0 && total_items > SIZE_MAX / item_size) {

PyErr_SetString(PyExc_OverflowError, "Tensor allocation too large");

return NULL;

}

size_t const total_bytes = total_items * item_size;

Tensor *tensor = PyObject_NewVar(Tensor, &TensorType, total_bytes + NK_TENSOR_PADDING_);

if (!tensor) {

PyErr_NoMemory();

return NULL;

}

tensor->dtype = dtype;

tensor->rank = rank;

for (size_t i = 0; i < NK_TENSOR_MAX_RANK; i++) {

tensor->shape[i] = (i < rank) ? shape[i] : 0;

tensor->strides[i] = 0;

}

if (rank > 0) {

tensor->strides[rank - 1] = (Py_ssize_t)item_size;

for (size_t i = rank - 1; i > 0; i--) tensor->strides[i - 1] = tensor->strides[i] * tensor->shape[i];

}

tensor->parent = NULL;

tensor->data = tensor->start;

return tensor;

}

Tensor *Tensor_view(Tensor *parent, char *data_ptr, nk_dtype_t dtype, size_t rank, Py_ssize_t const *shape,

Py_ssize_t const *strides) {

if (rank > NK_TENSOR_MAX_RANK) {

PyErr_Format(PyExc_ValueError, "View rank %zu exceeds maximum %d", rank, NK_TENSOR_MAX_RANK);

return NULL;

}

Tensor *view = PyObject_NewVar(Tensor, &TensorType, 0);

if (!view) {

PyErr_NoMemory();

return NULL;

}

view->dtype = dtype;

view->rank = rank;

for (size_t i = 0; i < NK_TENSOR_MAX_RANK; i++) {

view->shape[i] = (i < rank) ? shape[i] : 0;

view->strides[i] = (i < rank) ? strides[i] : 0;

}

view->parent = (PyObject *)parent;

Py_INCREF(parent);

view->data = data_ptr;

return view;

}

static Tensor *Tensor_scalar(nk_dtype_t dtype, void const *value) {

size_t const item_size = bytes_per_dtype(dtype);

Tensor *tensor = PyObject_NewVar(Tensor, &TensorType, item_size);

if (!tensor) {

PyErr_NoMemory();

return NULL;

}

tensor->dtype = dtype;

tensor->rank = 0;

for (size_t i = 0; i < NK_TENSOR_MAX_RANK; i++) {

tensor->shape[i] = 0;

tensor->strides[i] = 0;

}

tensor->parent = NULL;

tensor->data = tensor->start;

memcpy(tensor->data, value, item_size);

return tensor;

}

static PyObject *Tensor_float(PyObject *self) {

Tensor *tensor = (Tensor *)self;

if (tensor->rank != 0) {

PyErr_SetString(PyExc_TypeError, "only 0-dimensional tensors can be converted to float");

return NULL;

}

PyObject *scalar = tensor_read_scalar(tensor, 0);

if (!scalar) return NULL;

PyObject *result = PyNumber_Float(scalar);

Py_DECREF(scalar);

return result;

}

static PyObject *Tensor_int(PyObject *self) {

Tensor *tensor = (Tensor *)self;

if (tensor->rank != 0) {

PyErr_SetString(PyExc_TypeError, "only 0-dimensional tensors can be converted to int");

return NULL;

}

PyObject *scalar = tensor_read_scalar(tensor, 0);

if (!scalar) return NULL;

PyObject *result = PyNumber_Long(scalar);

Py_DECREF(scalar);

return result;

}

static PyObject *Tensor_positive(PyObject *self) { return Tensor_copy(self, NULL); }

void compute_contiguous_strides(size_t rank, Py_ssize_t const *shape, size_t item_size, Py_ssize_t *strides_out) {

if (rank == 0) return;

strides_out[rank - 1] = (Py_ssize_t)item_size;

for (size_t d = rank - 1; d > 0; --d) strides_out[d - 1] = strides_out[d] * shape[d];

}

static void linearize_cast_recursive( //

size_t remaining_dims, size_t contiguous_tail_dims) {

// Base case: all remaining dimensions are contiguous — one operation

if (remaining_dims <= contiguous_tail_dims) {

size_t slice_elements = 1;

for (size_t dim = 0; dim < remaining_dims; ++dim) slice_elements *= (size_t)shape[dim];

if (src_dtype == dest_dtype) memcpy(dest_data, src_data, slice_elements * src_element_size);

else nk_cast(src_data, src_dtype, (nk_size_t)slice_elements, dest_data, dest_dtype);

return;

}

// Recursive case: iterate outermost non-contiguous dimension

size_t const dim_extent = (size_t)shape[0];

// Compute the contiguous dest stride for this level

size_t inner_elements = 1;

for (size_t dim = 1; dim < remaining_dims; ++dim) inner_elements *= (size_t)shape[dim];

size_t const dest_row_bytes = inner_elements * dest_element_size;

for (size_t position = 0; position < dim_extent; ++position) {

linearize_cast_recursive( //

src_data + (Py_ssize_t)position * strides[0], src_dtype, //

dest_data + (Py_ssize_t)position * (Py_ssize_t)dest_row_bytes, //

dest_dtype, src_element_size, dest_element_size, //

shape + 1, strides + 1, //

remaining_dims - 1, contiguous_tail_dims);

}

}

void linearize_cast_into(char const *src_data, nk_dtype_t src_dtype, char *dest_data, nk_dtype_t dest_dtype,

size_t rank, Py_ssize_t const *shape, Py_ssize_t const *strides, size_t total_elements) {

nk_unused_(total_elements);

size_t src_element_size = bytes_per_dtype(src_dtype);

size_t dest_element_size = bytes_per_dtype(dest_dtype);

// Count how many trailing dims are contiguous in src

size_t contiguous_tail_dims = 0;

Py_ssize_t expected_stride = (Py_ssize_t)src_element_size;

for (size_t dim = rank; dim-- > 0;) {

if (strides[dim] != expected_stride) break;

expected_stride *= shape[dim];

contiguous_tail_dims++;

}

linearize_cast_recursive(src_data, src_dtype, dest_data, dest_dtype, src_element_size, dest_element_size, shape,

strides, rank, contiguous_tail_dims);

}

char *ensure_contiguous_buffer(char const *src_data, nk_dtype_t src_dtype, nk_dtype_t target_dtype, size_t rank,

Py_ssize_t const *shape, Py_ssize_t const *strides, size_t total_elements,

int *needs_free) {

size_t src_element_size = bytes_per_dtype(src_dtype);

size_t dest_element_size = bytes_per_dtype(target_dtype);

// Check full contiguity

int is_contiguous = 1;

Py_ssize_t expected_stride = (Py_ssize_t)src_element_size;

for (size_t dim = rank; dim-- > 0;) {

if (strides[dim] != expected_stride) {

is_contiguous = 0;

break;

}

expected_stride *= shape[dim];

}

// Zero-copy: contiguous + same dtype

if (is_contiguous && src_dtype == target_dtype) {

*needs_free = 0;

return (char *)src_data;

}

// Single allocation, delegate

char *output = PyMem_Malloc(total_elements * dest_element_size + NK_TENSOR_PADDING_);

if (!output) {

PyErr_NoMemory();

return NULL;

}

linearize_cast_into(src_data, src_dtype, output, target_dtype, rank, shape, strides, total_elements);

*needs_free = 1;

return output;

}

static PyObject *tensor_elementwise_scalar(Tensor *a, double alpha_value, double beta_value) {

nk_each_scale_punned_t kernel = NULL;

nk_capability_t cap = nk_cap_serial_k;

nk_find_kernel_punned(nk_kernel_each_scale_k, a->dtype, static_capabilities, (nk_kernel_punned_t *)&kernel, &cap);

if (!kernel || !cap) {

PyErr_Format(PyExc_NotImplementedError, "scale not supported for dtype '%s'", dtype_to_python_string(a->dtype));

return NULL;

}

Tensor *r = Tensor_new(a->dtype, a->rank, a->shape);

if (!r) return NULL;

size_t item_size = bytes_per_dtype(a->dtype);

Py_ssize_t r_strides[NK_TENSOR_MAX_RANK];

compute_contiguous_strides(a->rank, a->shape, item_size, r_strides);

Py_buffer a_buf = {.ndim = (int)a->rank,

.itemsize = (Py_ssize_t)item_size,

.shape = a->shape,

.strides = a->strides};

Py_buffer const *bufs[] = {&a_buf};

size_t contiguous_tail = shared_contiguous_tail_dimensions(bufs, 1, a->rank);

nk_scalar_buffer_t alpha_buf, beta_buf;

nk_dtype_t scalar_dtype = nk_each_scale_input_dtype(a->dtype);

nk_scalar_buffer_set_f64(&alpha_buf, alpha_value, scalar_dtype);

nk_scalar_buffer_set_f64(&beta_buf, beta_value, scalar_dtype);

PyThreadState *gil = PyEval_SaveThread();

each_scale_recursive(kernel, a->data, r->data, &alpha_buf, &beta_buf, a->shape, a->strides, r_strides, a->rank,

contiguous_tail);

PyEval_RestoreThread(gil);

return (PyObject *)r;

}

static PyObject *Tensor_negative(PyObject *self) { return tensor_elementwise_scalar((Tensor *)self, -1.0, 0.0); }

static PyObject *Tensor_add(PyObject *self, PyObject *other) {

if (!PyObject_TypeCheck(self, &TensorType)) { Py_RETURN_NOTIMPLEMENTED; }

Tensor *a = (Tensor *)self;

// Tensor([1,2,3]) + Tensor([4,5,6]) → element-wise sum via sum kernel

if (PyObject_TypeCheck(other, &TensorType)) {

Tensor *b = (Tensor *)other;

if (a->rank != b->rank || a->dtype != b->dtype) {

PyErr_SetString(PyExc_ValueError, "shape/dtype mismatch");

return NULL;

}

for (size_t i = 0; i < a->rank; i++)

if (a->shape[i] != b->shape[i]) {

PyErr_SetString(PyExc_ValueError, "shape mismatch");

return NULL;

}

nk_each_sum_punned_t kernel = NULL;

nk_capability_t cap = nk_cap_serial_k;

nk_find_kernel_punned(nk_kernel_each_sum_k, a->dtype, static_capabilities, (nk_kernel_punned_t *)&kernel, &cap);

if (!kernel || !cap) {

PyErr_Format(PyExc_NotImplementedError, "add not supported for dtype '%s'",

dtype_to_python_string(a->dtype));

return NULL;

}

Tensor *r = Tensor_new(a->dtype, a->rank, a->shape);

if (!r) return NULL;

size_t item_size = bytes_per_dtype(a->dtype);

Py_ssize_t r_strides[NK_TENSOR_MAX_RANK];

compute_contiguous_strides(a->rank, a->shape, item_size, r_strides);

Py_buffer a_buf = {.ndim = (int)a->rank,

.itemsize = (Py_ssize_t)item_size,

.shape = a->shape,

.strides = a->strides};

Py_buffer b_buf = {.ndim = (int)b->rank,

.itemsize = (Py_ssize_t)item_size,

.shape = b->shape,

.strides = b->strides};

Py_buffer const *bufs[] = {&a_buf, &b_buf};

size_t contiguous_tail = shared_contiguous_tail_dimensions(bufs, 2, a->rank);

PyThreadState *gil = PyEval_SaveThread();

each_sum_recursive(kernel, a->data, b->data, r->data, a->shape, a->strides, b->strides, r_strides, a->rank,

contiguous_tail);

PyEval_RestoreThread(gil);

return (PyObject *)r;

}

// Tensor([1,2,3]) + 5.0 → broadcast scalar addition via scale kernel (α=1, β=scalar)

if (PyFloat_Check(other) || PyLong_Check(other)) {

double sc = PyFloat_Check(other) ? PyFloat_AsDouble(other) : (double)PyLong_AsLong(other);

return tensor_elementwise_scalar(a, 1.0, sc);

}

Py_RETURN_NOTIMPLEMENTED;

}

static PyObject *Tensor_subtract(PyObject *self, PyObject *other) {

if (!PyObject_TypeCheck(self, &TensorType)) { Py_RETURN_NOTIMPLEMENTED; }

Tensor *a = (Tensor *)self;

// Tensor([4,5,6]) - Tensor([1,2,3]) → element-wise difference via blend kernel (α=1, β=−1)

if (PyObject_TypeCheck(other, &TensorType)) {

Tensor *b = (Tensor *)other;

if (a->rank != b->rank || a->dtype != b->dtype) {

PyErr_SetString(PyExc_ValueError, "shape/dtype mismatch");

return NULL;

}

for (size_t i = 0; i < a->rank; i++)

if (a->shape[i] != b->shape[i]) {

PyErr_SetString(PyExc_ValueError, "shape mismatch");

return NULL;

}

// Single-pass subtract via blend: result = 1*a + (-1)*b

nk_each_blend_punned_t kernel = NULL;

nk_capability_t cap = nk_cap_serial_k;

nk_find_kernel_punned(nk_kernel_each_blend_k, a->dtype, static_capabilities, (nk_kernel_punned_t *)&kernel,

&cap);

if (!kernel || !cap) {

PyErr_Format(PyExc_NotImplementedError, "subtract not supported for dtype '%s'",

dtype_to_python_string(a->dtype));

return NULL;

}

Tensor *r = Tensor_new(a->dtype, a->rank, a->shape);

if (!r) return NULL;

size_t item_size = bytes_per_dtype(a->dtype);

Py_ssize_t r_strides[NK_TENSOR_MAX_RANK];

compute_contiguous_strides(a->rank, a->shape, item_size, r_strides);

Py_buffer a_buf = {.ndim = (int)a->rank,

.itemsize = (Py_ssize_t)item_size,

.shape = a->shape,

.strides = a->strides};

Py_buffer b_buf = {.ndim = (int)b->rank,

.itemsize = (Py_ssize_t)item_size,

.shape = b->shape,

.strides = b->strides};

Py_buffer const *bufs[] = {&a_buf, &b_buf};

size_t contiguous_tail = shared_contiguous_tail_dimensions(bufs, 2, a->rank);

nk_scalar_buffer_t alpha_buf, beta_buf;

nk_dtype_t scalar_dtype = nk_each_scale_input_dtype(a->dtype);

nk_scalar_buffer_set_f64(&alpha_buf, 1.0, scalar_dtype);

nk_scalar_buffer_set_f64(&beta_buf, -1.0, scalar_dtype);

PyThreadState *gil = PyEval_SaveThread();

each_blend_recursive(kernel, a->data, b->data, r->data, &alpha_buf, &beta_buf, a->shape, a->strides, b->strides,

r_strides, a->rank, contiguous_tail);

PyEval_RestoreThread(gil);

return (PyObject *)r;

}

// Tensor([4,5,6]) - 1.0 → broadcast scalar subtraction via scale kernel (α=1, β=−scalar)

if (PyFloat_Check(other) || PyLong_Check(other)) {

double sc = PyFloat_Check(other) ? PyFloat_AsDouble(other) : (double)PyLong_AsLong(other);

return tensor_elementwise_scalar(a, 1.0, -sc);

}

Py_RETURN_NOTIMPLEMENTED;

}

static PyObject *Tensor_multiply(PyObject *self, PyObject *other) {

if (!PyObject_TypeCheck(self, &TensorType)) { Py_RETURN_NOTIMPLEMENTED; }

Tensor *a = (Tensor *)self;

// Tensor([1,2,3]) × Tensor([4,5,6]) → element-wise product via fma kernel (α=1, β=0)

if (PyObject_TypeCheck(other, &TensorType)) {

Tensor *b = (Tensor *)other;

if (a->rank != b->rank || a->dtype != b->dtype) {

PyErr_SetString(PyExc_ValueError, "shape/dtype mismatch");

return NULL;

}

for (size_t i = 0; i < a->rank; i++)

if (a->shape[i] != b->shape[i]) {

PyErr_SetString(PyExc_ValueError, "shape mismatch");

return NULL;

}

nk_each_fma_punned_t kernel = NULL;

nk_capability_t cap = nk_cap_serial_k;

nk_find_kernel_punned(nk_kernel_each_fma_k, a->dtype, static_capabilities, (nk_kernel_punned_t *)&kernel, &cap);

if (!kernel || !cap) {

PyErr_Format(PyExc_NotImplementedError, "multiply not supported for dtype '%s'",

dtype_to_python_string(a->dtype));

return NULL;

}

Tensor *r = Tensor_new(a->dtype, a->rank, a->shape);

if (!r) return NULL;

size_t item_size = bytes_per_dtype(a->dtype);

size_t total_items = 1;

for (size_t i = 0; i < a->rank; i++) total_items *= (size_t)a->shape[i];

memset(r->data, 0, total_items * item_size); // prevent 0*NaN=NaN from uninitialized memory

Py_ssize_t r_strides[NK_TENSOR_MAX_RANK];

compute_contiguous_strides(a->rank, a->shape, item_size, r_strides);

Py_buffer a_buf = {.ndim = (int)a->rank,

.itemsize = (Py_ssize_t)item_size,

.shape = a->shape,

.strides = a->strides};

Py_buffer b_buf = {.ndim = (int)b->rank,

.itemsize = (Py_ssize_t)item_size,

.shape = b->shape,

.strides = b->strides};

Py_buffer const *bufs[] = {&a_buf, &b_buf};

size_t contiguous_tail = shared_contiguous_tail_dimensions(bufs, 2, a->rank);

// fma(a, b, dummy, n, alpha=1, beta=0) -> 1*a*b + 0*dummy

nk_scalar_buffer_t alpha_buf, beta_buf;

nk_dtype_t scalar_dtype = nk_each_scale_input_dtype(a->dtype);

nk_scalar_buffer_set_f64(&alpha_buf, 1.0, scalar_dtype);

nk_scalar_buffer_set_f64(&beta_buf, 0.0, scalar_dtype);

PyThreadState *gil = PyEval_SaveThread();

each_fma_recursive(kernel, a->data, b->data, r->data, r->data, &alpha_buf, &beta_buf, a->shape, a->strides,

b->strides, r_strides, r_strides, a->rank, contiguous_tail);

PyEval_RestoreThread(gil);

return (PyObject *)r;

}

// Tensor([1,2,3]) × 5.0 → broadcast scalar multiply via scale kernel (α=scalar, β=0)

if (PyFloat_Check(other) || PyLong_Check(other)) {

double sc = PyFloat_Check(other) ? PyFloat_AsDouble(other) : (double)PyLong_AsLong(other);

return tensor_elementwise_scalar(a, sc, 0.0);

}

Py_RETURN_NOTIMPLEMENTED;

}

static PyNumberMethods Tensor_as_number = {

.nb_add = Tensor_add,

.nb_subtract = Tensor_subtract,

.nb_multiply = Tensor_multiply,

.nb_matrix_multiply = Tensor_matmul,

.nb_negative = Tensor_negative,

.nb_positive = Tensor_positive,

.nb_float = Tensor_float,

.nb_int = Tensor_int,

};

static PyObject *Tensor_get_shape(PyObject *self, void *closure) {

nk_unused_(closure);

Tensor *tensor = (Tensor *)self;

PyObject *shape_tuple = PyTuple_New(tensor->rank);

if (!shape_tuple) return NULL;

for (size_t i = 0; i < tensor->rank; i++) {

PyTuple_SET_ITEM(shape_tuple, i, PyLong_FromSsize_t(tensor->shape[i]));

}

return shape_tuple;

}

static PyObject *Tensor_get_dtype(PyObject *self, void *closure) {

nk_unused_(closure);

Tensor *tensor = (Tensor *)self;

return PyUnicode_FromString(dtype_to_string(tensor->dtype));

}

static PyObject *Tensor_get_ndim(PyObject *self, void *closure) {

nk_unused_(closure);

Tensor *tensor = (Tensor *)self;

return PyLong_FromSize_t(tensor->rank);

}

static PyObject *Tensor_get_size(PyObject *self, void *closure) {

nk_unused_(closure);

Tensor *tensor = (Tensor *)self;

Py_ssize_t total = 1;

for (size_t i = 0; i < tensor->rank; i++) total *= tensor->shape[i];

return PyLong_FromSsize_t(total);

}

static PyObject *Tensor_get_nbytes(PyObject *self, void *closure) {

nk_unused_(closure);

Tensor *tensor = (Tensor *)self;

Py_ssize_t total = 1;

for (size_t i = 0; i < tensor->rank; i++) total *= tensor->shape[i];

return PyLong_FromSsize_t(total * (Py_ssize_t)bytes_per_dtype(tensor->dtype));

}

static PyObject *Tensor_get_strides(PyObject *self, void *closure) {

nk_unused_(closure);

Tensor *tensor = (Tensor *)self;

PyObject *strides_tuple = PyTuple_New(tensor->rank);

if (!strides_tuple) return NULL;

for (size_t i = 0; i < tensor->rank; i++) {

PyTuple_SET_ITEM(strides_tuple, i, PyLong_FromSsize_t(tensor->strides[i]));

}

return strides_tuple;

}

static PyObject *Tensor_get_itemsize(PyObject *self, void *closure) {

nk_unused_(closure);

Tensor *tensor = (Tensor *)self;

return PyLong_FromSize_t(bytes_per_dtype(tensor->dtype));

}

static PyObject *Tensor_get_T(PyObject *self, void *closure) {

nk_unused_(closure);

Tensor *tensor = (Tensor *)self;

if (tensor->rank < 2) {

// 0D or 1D: transpose is a view of itself

Py_INCREF(self);

return self;

}

// Reverse shape and strides

Py_ssize_t new_shape[NK_TENSOR_MAX_RANK];

Py_ssize_t new_strides[NK_TENSOR_MAX_RANK];

for (size_t i = 0; i < tensor->rank; i++) {

new_shape[i] = tensor->shape[tensor->rank - 1 - i];

new_strides[i] = tensor->strides[tensor->rank - 1 - i];

}

Tensor *root_parent = tensor->parent ? (Tensor *)tensor->parent : tensor;

return (PyObject *)Tensor_view(root_parent, tensor->data, tensor->dtype, tensor->rank, new_shape, new_strides);

}

static PyObject *Tensor_get_array_interface(PyObject *self, void *closure) {

nk_unused_(closure);

Tensor *tensor = (Tensor *)self;

PyObject *dict = PyDict_New();

if (!dict) return NULL;

PyObject *shape = Tensor_get_shape(self, NULL);

if (!shape) {

Py_DECREF(dict);

return NULL;

}

PyDict_SetItemString(dict, "shape", shape);

Py_DECREF(shape);

char const *typestr = dtype_to_array_typestr(tensor->dtype);

PyObject *typestr_obj = PyUnicode_FromString(typestr);

if (!typestr_obj) {

Py_DECREF(dict);

return NULL;

}

PyDict_SetItemString(dict, "typestr", typestr_obj);

Py_DECREF(typestr_obj);

PyObject *data_ptr = PyLong_FromVoidPtr(tensor->data);

if (!data_ptr) {

Py_DECREF(dict);

return NULL;

}

PyObject *data_tuple = PyTuple_Pack(2, data_ptr, Py_False);

Py_DECREF(data_ptr);

if (!data_tuple) {

Py_DECREF(dict);

return NULL;

}

PyDict_SetItemString(dict, "data", data_tuple);

Py_DECREF(data_tuple);

PyObject *strides = Tensor_get_strides(self, NULL);

if (!strides) {

Py_DECREF(dict);

return NULL;

}

PyDict_SetItemString(dict, "strides", strides);

Py_DECREF(strides);

PyObject *version = PyLong_FromLong(3);

if (!version) {

Py_DECREF(dict);

return NULL;

}

PyDict_SetItemString(dict, "version", version);

Py_DECREF(version);

return dict;

}

static PyObject *Tensor_get_is_contiguous(PyObject *self, void *closure) {

nk_unused_(closure);

Tensor *tensor = (Tensor *)self;

size_t item_size = bytes_per_dtype(tensor->dtype);

return PyBool_FromLong(tensor_is_c_contig(tensor, item_size));

}

static PyGetSetDef Tensor_getset[] = {

{"shape", Tensor_get_shape, NULL, "Shape of the array", NULL},

{"dtype", Tensor_get_dtype, NULL, "Data type of the array", NULL},

{"ndim", Tensor_get_ndim, NULL, "Number of dimensions", NULL},

{"size", Tensor_get_size, NULL, "Total number of elements", NULL},

{"nbytes", Tensor_get_nbytes, NULL, "Total bytes of data", NULL},

{"strides", Tensor_get_strides, NULL, "Strides in bytes", NULL},

{"itemsize", Tensor_get_itemsize, NULL, "Size of one element in bytes", NULL},

{"T", Tensor_get_T, NULL, "Transposed view of the array", NULL},

{"__array_interface__", Tensor_get_array_interface, NULL, "NumPy array interface", NULL},

{"is_contiguous", Tensor_get_is_contiguous, NULL, "Whether the tensor is C-contiguous", NULL},

{NULL, NULL, NULL, NULL, NULL},

};

PyObject *Tensor_copy(PyObject *self, PyObject *args) {

nk_unused_(args);

Tensor *tensor = (Tensor *)self;

size_t total_elements = 1;

for (size_t i = 0; i < tensor->rank; i++) total_elements *= (size_t)tensor->shape[i];

Tensor *result = Tensor_new(tensor->dtype, tensor->rank, tensor->shape);

if (!result) return NULL;

linearize_cast_into(tensor->data, tensor->dtype, result->data, tensor->dtype, tensor->rank, tensor->shape,

tensor->strides, total_elements);

return (PyObject *)result;

}

PyObject *Tensor_reshape(PyObject *self, PyObject *const *args, Py_ssize_t nargs) {

Tensor *tensor = (Tensor *)self;

Py_ssize_t new_shape[NK_TENSOR_MAX_RANK];

size_t new_rank = 0;

if (nargs == 1 && PyTuple_Check(args[0])) {

PyObject *shape_tuple = args[0];

new_rank = PyTuple_GET_SIZE(shape_tuple);

if (new_rank > NK_TENSOR_MAX_RANK) {

PyErr_Format(PyExc_ValueError, "reshape: too many dimensions (%zu > %d)", new_rank, NK_TENSOR_MAX_RANK);

return NULL;

}

for (size_t i = 0; i < new_rank; i++) {

PyObject *item = PyTuple_GET_ITEM(shape_tuple, i);

if (!PyLong_Check(item)) {

PyErr_SetString(PyExc_TypeError, "reshape: shape dimensions must be integers");

return NULL;

}

new_shape[i] = PyLong_AsSsize_t(item);

if (new_shape[i] < 0) {

PyErr_SetString(PyExc_ValueError, "reshape: negative dimensions not supported");

return NULL;

}

}

}

else {

new_rank = (size_t)nargs;

if (new_rank > NK_TENSOR_MAX_RANK) {

PyErr_Format(PyExc_ValueError, "reshape: too many dimensions (%zu > %d)", new_rank, NK_TENSOR_MAX_RANK);

return NULL;

}

for (size_t i = 0; i < new_rank; i++) {

PyObject *item = args[i];

if (!PyLong_Check(item)) {

PyErr_SetString(PyExc_TypeError, "reshape: shape dimensions must be integers");

return NULL;

}

new_shape[i] = PyLong_AsSsize_t(item);

if (new_shape[i] < 0) {

PyErr_SetString(PyExc_ValueError, "reshape: negative dimensions not supported");

return NULL;

}

}

}

Py_ssize_t new_total = 1;

for (size_t i = 0; i < new_rank; i++) new_total *= new_shape[i];

Py_ssize_t old_total = 1;

for (size_t i = 0; i < tensor->rank; i++) old_total *= tensor->shape[i];

if (new_total != old_total) {

PyErr_Format(PyExc_ValueError, "reshape: cannot reshape tensor of size %zd into shape with size %zd", old_total,

new_total);

return NULL;

}

size_t item_size = bytes_per_dtype(tensor->dtype);

if (tensor_is_c_contig(tensor, item_size)) {

Py_ssize_t new_strides[NK_TENSOR_MAX_RANK];

Py_ssize_t stride = item_size;

for (size_t i = new_rank; i > 0; i--) {

new_strides[i - 1] = stride;

stride *= new_shape[i - 1];

}

Tensor *root_parent = tensor->parent ? (Tensor *)tensor->parent : tensor;

return (PyObject *)Tensor_view(root_parent, tensor->data, tensor->dtype, new_rank, new_shape, new_strides);