// Copyright 2026 The Tenzr Project Authors
//
// SPDX-License-Identifier: Apache-2.0
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// https://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
#include <cctype>
#include <cstring>
#include <functional>
#include <iostream>
#include <sstream>
#include "bindings.h"
namespace tenzr::python {
template <typename Nd> Array ArrayFromNdarray(const Nd& nd, std::optional<DType> dtype_opt) {
size_t rank = nd.ndim();
std::vector<int64_t> dims(rank);
for (size_t i = 0; i < rank; ++i) {
dims[i] = static_cast<int64_t>(nd.shape(i));
}
Shape shape(dims);
DType src_dt = FromDlpackDtype(nd.dtype());
if (shape.NumElements() == 1 && txpr::TracerContext::IsActive()) {
Array scalar = Array::Full(shape, ReadScalarAsDouble(nd.data(), src_dt), src_dt);
return (dtype_opt.has_value() && *dtype_opt != src_dt) ? scalar.Astype(*dtype_opt) : scalar;
}
Array arr = Array::Parameter(shape, src_dt);
size_t total_bytes = static_cast<size_t>(shape.NumElements()) * SizeOf(src_dt);
auto buf = MemoryPool::Instance().Allocate(total_bytes);
std::memcpy(buf->Data(), nd.data(), total_bytes);
arr.BindData(buf);
return (dtype_opt.has_value() && *dtype_opt != src_dt) ? arr.Astype(*dtype_opt) : arr;
}
// Map an ml_dtypes (non-standard) numpy dtype name to a Tenzr DType.
bool TryMlDtypeFromName(const std::string& name, DType& out) {
static const std::unordered_map<std::string, DType> kMap = {
{"bfloat16", DType::kBFloat16},
{"float8_e3m4", DType::kFloat8E3M4},
{"float8_e4m3", DType::kFloat8E4M3},
{"float8_e4m3fn", DType::kFloat8E4M3FN},
{"float8_e4m3fnuz", DType::kFloat8E4M3FNUZ},
{"float8_e4m3b11fnuz", DType::kFloat8E4M3B11FNUZ},
{"float8_e5m2", DType::kFloat8E5M2},
{"float8_e5m2fnuz", DType::kFloat8E5M2FNUZ},
{"float8_e8m0fnu", DType::kFloat8E8M0FNU},
{"float6_e2m3fn", DType::kFloat6E2M3FN},
{"float6_e3m2fn", DType::kFloat6E3M2FN},
{"float4_e2m1fn", DType::kFloat4E2M1FN},
{"int1", DType::kInt1},
{"uint1", DType::kUInt1},
{"int2", DType::kInt2},
{"uint2", DType::kUInt2},
{"int4", DType::kInt4},
{"uint4", DType::kUInt4},
};
auto it = kMap.find(name);
if (it == kMap.end()) {
return false;
}
out = it->second;
return true;
}
// Helper to convert any Python object (list, float, int, ndarray) to tenzr::Array
Array PyToArray(nb::handle obj, std::optional<DType> dtype_opt = std::nullopt) {
if (nb::isinstance<Array>(obj)) {
Array arr = nb::cast<Array>(obj);
if (dtype_opt.has_value() && arr.Dtype() != *dtype_opt) {
return plax::ConvertElementType(arr, *dtype_opt);
}
return arr;
}
// Fast path 1: Python scalars, without importing NumPy.
if (nb::isinstance<nb::bool_>(obj) || nb::isinstance<nb::int_>(obj) || nb::isinstance<nb::float_>(obj)) {
DType dt = dtype_opt.value_or(DefaultScalarDtype(obj));
return Array::Full(Shape({}), nb::cast<double>(obj), dt);
}
// Fast path 2: any object exposing the Python buffer protocol, without importing NumPy.
const bool is_text = nb::isinstance<nb::str>(obj) || nb::isinstance<nb::bytes>(obj);
if (!is_text && obj.is_valid() && PyObject_CheckBuffer(obj.ptr())) {
try {
auto nd = nb::cast<nb::ndarray<nb::ro, nb::c_contig>>(obj);
return ArrayFromNdarray(nd, dtype_opt);
} catch (const nb::cast_error&) {
// Non-contiguous or otherwise unsupported buffer: fall through to NumPy.
}
}
// Fallback: NumPy handles arbitrary sequences (lists/tuples) and exotic objects.
nb::object np = nb::module_::import_("numpy");
nb::object nparr;
try {
nparr = np.attr("ascontiguousarray")(np.attr("asarray")(obj));
} catch (const nb::python_error&) {
throw nb::type_error("cannot convert object to a tenzr Array");
}
try {
auto nd = nb::cast<nb::ndarray<nb::ro, nb::c_contig>>(nparr);
return ArrayFromNdarray(nd, dtype_opt);
} catch (const nb::cast_error&) {
// ml_dtypes / non-standard dtype: nanobind rejects the opaque buffer, so
// map the numpy dtype name and copy the raw bytes.
std::string dname = nb::cast<std::string>(nparr.attr("dtype").attr("name"));
DType src_dt;
if (TryMlDtypeFromName(dname, src_dt)) {
std::vector<int64_t> dims;
for (nb::handle h : nparr.attr("shape")) {
dims.push_back(nb::cast<int64_t>(h));
}
Shape shape(dims);
nb::bytes raw = nb::cast<nb::bytes>(nparr.attr("tobytes")());
const size_t nbytes = static_cast<size_t>(shape.NumElements()) * SizeOf(src_dt);
auto buf = MemoryPool::Instance().Allocate(nbytes);
std::memcpy(buf->Data(), raw.data(), nbytes);
Array arr = Array::Parameter(shape, src_dt);
arr.BindData(buf);
return (dtype_opt.has_value() && *dtype_opt != src_dt) ? arr.Astype(*dtype_opt) : arr;
}
throw nb::type_error(("cannot convert object to a tenzr Array (unsupported dtype '" + dname + "')").c_str());
}
}
// Convert tenzr::Array to nb::ndarray (honors StridedLayout offset + strides).
nb::object ArrayToNdarray(Array a, [[maybe_unused]] nb::handle owner) {
if (!a.HasData() && !a.IsSymbolic()) {
a = a.Eval();
}
TENZR_CHECK(a.HasData(), "Cannot export an unevaluated symbolic array to Python");
a.SyncHost();
TENZR_CHECK(a.Data() != nullptr, "Array has no host buffer to export");
const StridedLayout& lay = a.Layout();
size_t rank = lay.shape().Rank();
std::vector<size_t> shape_vec(rank);
std::vector<int64_t> stride_vec(rank);
for (size_t i = 0; i < rank; ++i) {
shape_vec[i] = static_cast<size_t>(lay.shape().Dim(i));
stride_vec[i] = lay.Stride(i);
}
size_t elem_size = SizeOf(a.Dtype());
void* ptr = a.Data()->Data();
uint64_t byte_offset = static_cast<uint64_t>(lay.Offset()) * static_cast<uint64_t>(elem_size);
auto dl_dtype = ToDlpackDtype(a.Dtype());
nb::capsule OwnerCapsule(new Array(a), [](void* p) noexcept { delete static_cast<Array*>(p); });
nb::ndarray<nb::numpy> nd(ptr, rank, shape_vec.data(), OwnerCapsule,
rank == 0 ? nullptr : stride_vec.data(), // element strides (DLPack)
dl_dtype,
1, // kDLCPU
0, 'C', byte_offset);
nb::object result = nb::cast(nd);
// numpy has no native dtype for bfloat16/narrow/sub-byte types, so nanobind
// silently produces an object array - a correctness hazard. Fail loudly.
if (nb::cast<std::string>(result.attr("dtype").attr("kind")) == "O") {
TENZR_THROW("cannot export " + std::string(ToString(a.Dtype())) +
" to numpy (no native numpy dtype); use tolist()/item(), .astype() to a standard dtype, or ml_dtypes");
}
return result;
}
// JAX-style weak-type promotion for Python scalar/array operands (D6).
inline DType PromotedBinaryDtype(const Array& a, nb::handle other, bool force_float) {
DType other_dt = nb::isinstance<Array>(other) ? nb::cast<Array>(other).Dtype() : PythonScalarDtype(other);
DType p = PromoteTypes(a.Dtype(), other_dt);
if (force_float && IsIntegral(p)) {
p = DType::kFloat32;
}
return p;
}
template <typename Op> inline Array BinaryArith(const Array& a, nb::handle other, bool force_float, Op op) {
DType p = PromotedBinaryDtype(a, other, force_float);
Array lhs = (a.Dtype() == p) ? a : a.Astype(p);
Array rhs = PyToArray(other, p);
return op(lhs, rhs);
}
template <typename Cmp> inline Array BinaryCmp(const Array& a, nb::handle other, Cmp cmp) {
DType p = PromotedBinaryDtype(a, other, false);
Array lhs = (a.Dtype() == p) ? a : a.Astype(p);
Array rhs = PyToArray(other, p);
return cmp(lhs, rhs);
}
// Basic indexing: integer, slice, ellipsis, and tuples thereof (D3).
// Never silently returns the whole array; unsupported keys raise.
inline Array ArrayGetitem(const Array& a_in, nb::handle key) {
const int64_t rank = static_cast<int64_t>(a_in.shape().Rank());
const bool symbolic = a_in.IsSymbolic();
const bool grad_active = txpr::IsGradTracingActive();
Array a = a_in;
if (!symbolic && !grad_active && !a_in.HasData()) {
a = a_in.Eval();
}
std::vector<nb::handle> comps;
if (nb::isinstance<nb::tuple>(key)) {
nb::tuple t = nb::cast<nb::tuple>(key);
for (size_t i = 0; i < t.size(); ++i) {
comps.push_back(t[i]);
}
} else {
comps.push_back(key);
}
for (auto comp : comps) {
const bool supported = nb::isinstance<nb::slice>(comp) || nb::isinstance<nb::int_>(comp) || nb::isinstance<nb::ellipsis>(comp);
if (!supported) {
throw nb::type_error(
"tenzr.Array indexing supports integers, slices, Ellipsis, and tuples thereof; "
"advanced (array/list/boolean) indexing and newaxis are not yet supported");
}
}
int64_t indexed = 0;
for (auto comp : comps) {
if (!nb::isinstance<nb::ellipsis>(comp)) {
indexed++;
}
}
TENZR_CHECK(indexed <= rank, "too many indices for array");
const int64_t ellipsis_dims = rank - indexed;
std::vector<int64_t> src_strides(static_cast<size_t>(rank));
int64_t src_offset = 0;
if (symbolic) {
int64_t s = 1;
for (int64_t i = rank - 1; i >= 0; --i) {
src_strides[static_cast<size_t>(i)] = s;
s *= a.shape().Dim(static_cast<size_t>(i));
}
} else {
const auto& lay = a.Layout();
for (int64_t i = 0; i < rank; ++i) {
src_strides[static_cast<size_t>(i)] = lay.Stride(static_cast<size_t>(i));
}
src_offset = lay.Offset();
}
std::vector<Slice> slices;
slices.reserve(static_cast<size_t>(rank));
std::vector<int64_t> out_dims;
std::vector<int64_t> out_strides;
int64_t out_offset = src_offset;
std::vector<int64_t> squeeze_axes;
int64_t dim = 0;
auto emit_dim = [&](int64_t start, int64_t step, int64_t len) {
const int64_t strd = src_strides[static_cast<size_t>(dim)];
out_dims.push_back(len);
out_strides.push_back(strd * step);
out_offset += start * strd;
slices.push_back(Slice(start, start + len * step, step));
dim++;
};
auto emit_full = [&]() {
out_dims.push_back(a.shape().Dim(static_cast<size_t>(dim)));
out_strides.push_back(src_strides[static_cast<size_t>(dim)]);
slices.push_back(Slice::All());
dim++;
};
for (auto comp : comps) {
if (nb::isinstance<nb::ellipsis>(comp)) {
for (int64_t k = 0; k < ellipsis_dims; ++k) {
emit_full();
}
} else if (nb::isinstance<nb::int_>(comp)) {
int64_t idx = nb::cast<int64_t>(comp);
const int64_t d = a.shape().Dim(static_cast<size_t>(dim));
if (idx < 0) {
idx += d;
}
if (idx < 0 || idx >= d) {
throw nb::index_error("index out of bounds");
}
const int64_t strd = src_strides[static_cast<size_t>(dim)];
out_offset += idx * strd;
slices.push_back(Slice(idx, idx + 1, 1));
squeeze_axes.push_back(static_cast<int64_t>(out_dims.size()));
out_dims.push_back(1);
out_strides.push_back(strd);
dim++;
} else // slice
{
auto s = nb::cast<nb::slice>(comp);
const int64_t d = a.shape().Dim(static_cast<size_t>(dim));
auto [start, stop, step, len] = s.compute(static_cast<size_t>(d));
(void)stop;
if (step < 0 && (symbolic || grad_active)) {
TENZR_THROW("negative-step slicing is not supported inside traced/differentiated functions");
}
emit_dim(static_cast<int64_t>(start), static_cast<int64_t>(step), static_cast<int64_t>(len));
}
}
while (dim < rank) {
emit_full();
}
Array result = (symbolic || grad_active) ? a.slice(slices) : a.AsStrided(Shape(out_dims), out_strides, out_offset);
for (auto it = squeeze_axes.rbegin(); it != squeeze_axes.rend(); ++it) {
result = result.Squeeze(*it);
}
return result;
}
// Normalize a Python axis argument (None / int / sequence of ints) to a vector.
inline std::vector<int64_t> ParseAxes(nb::handle axis) {
std::vector<int64_t> axes;
if (axis.is_none()) {
return axes;
}
if (nb::isinstance<nb::int_>(axis)) {
axes.push_back(nb::cast<int64_t>(axis));
return axes;
}
if (nb::isinstance<nb::tuple>(axis) || nb::isinstance<nb::list>(axis)) {
for (auto item : nb::cast<nb::sequence>(axis)) {
axes.push_back(nb::cast<int64_t>(item));
}
return axes;
}
throw nb::type_error("axis must be an int, a sequence of ints, or None");
}
// ---- Minimal DLPack C ABI (v0.8 and v1.0) for zero-copy consumption ----
namespace dlpack_abi {
constexpr int32_t kDLCPU = 1;
struct DLDevice {
int32_t device_type;
int32_t device_id;
};
struct DLDataType {
uint8_t code;
uint8_t bits;
uint16_t lanes;
};
struct DLTensor {
void* data;
DLDevice device;
int32_t ndim;
DLDataType dtype;
int64_t* shape;
int64_t* strides;
uint64_t byte_offset;
};
struct DLManagedTensor {
DLTensor dl_tensor;
void* manager_ctx;
void (*deleter)(DLManagedTensor*);
};
struct DLPackVersion {
uint32_t major;
uint32_t minor;
};
struct DLManagedTensorVersioned {
DLPackVersion version;
void* manager_ctx;
void (*deleter)(DLManagedTensorVersioned*);
uint64_t flags;
DLTensor dl_tensor;
};
} // namespace dlpack_abi
// ---- Device (GPU) DLPack export & import ----
namespace {
constexpr int32_t kDLCUDA_DEV = 2;
constexpr int32_t kDLROCM_DEV = 10;
constexpr int32_t kDLOneAPI_DEV = 14;
std::string LowerPlatform() {
auto& hw = HardwareManager::Instance();
if (!hw.IsInitialized()) {
try {
hw.Initialize();
} catch (const std::exception&) {
return std::string();
}
}
const std::string& p = hw.PlatformName();
std::string lp;
lp.reserve(p.size());
for (char c : p) {
lp.push_back(static_cast<char>(std::tolower(static_cast<unsigned char>(c))));
}
return lp;
}
int32_t DlDeviceTypeForPlatform() {
const std::string lp = LowerPlatform();
if (lp.find("rocm") != std::string::npos || lp.find("hip") != std::string::npos) {
return kDLROCM_DEV;
}
if (lp.find("oneapi") != std::string::npos || lp.find("xpu") != std::string::npos) {
return kDLOneAPI_DEV;
}
if (lp.find("cuda") != std::string::npos || lp.find("gpu") != std::string::npos) {
return kDLCUDA_DEV;
}
return 1; // kDLCPU
}
bool PlatformIsGpu() { return DlDeviceTypeForPlatform() != 1; }
const PJRT_Api* PjrtApi() {
auto& hw = HardwareManager::Instance();
if (!hw.IsInitialized()) {
try {
hw.Initialize();
} catch (const std::exception&) {
return nullptr;
}
}
return hw.Api();
}
bool DeviceExportSupported() {
const PJRT_Api* api = PjrtApi();
return api != nullptr && api->PJRT_Buffer_UnsafePointer != nullptr && api->PJRT_Buffer_IncreaseExternalReferenceCount != nullptr &&
api->PJRT_Buffer_DecreaseExternalReferenceCount != nullptr;
}
bool DeviceImportSupported() {
const PJRT_Api* api = PjrtApi();
return api != nullptr && HardwareManager::Instance().Client() != nullptr && api->PJRT_Client_CreateViewOfDeviceBuffer != nullptr;
}
struct DeviceExportCtx {
std::shared_ptr<DeviceBuffer> buffer;
int64_t* shape{nullptr};
};
void DeviceTensorDeleter(dlpack_abi::DLManagedTensor* self) {
auto* ctx = static_cast<DeviceExportCtx*>(self->manager_ctx);
if (ctx != nullptr) {
if (ctx->buffer) {
ctx->buffer->DecreaseExternalReferenceCount();
}
delete[] ctx->shape;
delete ctx;
}
delete self;
}
void DeviceCapsuleDestructor(PyObject* capsule) {
const char* name = PyCapsule_GetName(capsule);
if (name != nullptr && std::strcmp(name, "dltensor") == 0) {
auto* mt = static_cast<dlpack_abi::DLManagedTensor*>(PyCapsule_GetPointer(capsule, name));
if (mt != nullptr && mt->deleter != nullptr) {
mt->deleter(mt);
}
}
}
nb::object MakeDeviceDlpack(const Array& arr) {
Array ready = arr.HasDeviceData() ? arr : arr.Eval();
TENZR_CHECK(ready.HasDeviceData(), "Array has no device buffer to export");
std::shared_ptr<DeviceBuffer> db = ready.DeviceData();
TENZR_CHECK(db != nullptr, "Array device buffer is null");
db->AwaitReady();
db->IncreaseExternalReferenceCount();
void* ptr = db->UnsafePointer();
auto* mt = new dlpack_abi::DLManagedTensor();
std::memset(mt, 0, sizeof(*mt));
mt->dl_tensor.data = ptr;
mt->dl_tensor.device.device_type = DlDeviceTypeForPlatform();
mt->dl_tensor.device.device_id = db->DeviceId();
nb::dlpack::dtype dl_dt = ToDlpackDtype(ready.Dtype());
mt->dl_tensor.dtype.code = dl_dt.code;
mt->dl_tensor.dtype.bits = dl_dt.bits;
mt->dl_tensor.dtype.lanes = dl_dt.lanes;
mt->dl_tensor.ndim = static_cast<int32_t>(ready.shape().Rank());
const size_t rank = ready.shape().Rank();
auto* shape = new int64_t[rank == 0 ? 1 : rank];
for (size_t i = 0; i < rank; ++i) {
shape[i] = ready.shape().Dim(i);
}
mt->dl_tensor.shape = shape;
mt->dl_tensor.strides = nullptr;
mt->dl_tensor.byte_offset = 0;
mt->manager_ctx = new DeviceExportCtx{db, shape};
mt->deleter = &DeviceTensorDeleter;
PyObject* cap = PyCapsule_New(mt, "dltensor", &DeviceCapsuleDestructor);
if (cap == nullptr) {
mt->deleter(mt);
throw nb::python_error();
}
return nb::steal<nb::object>(cap);
}
struct DlpackImportCtx {
bool versioned;
void* managed;
};
void ImportViewDeleter(void* /*device_buffer_ptr*/, void* user_arg) {
auto* ctx = static_cast<DlpackImportCtx*>(user_arg);
if (ctx != nullptr) {
if (ctx->versioned) {
auto* mt = static_cast<dlpack_abi::DLManagedTensorVersioned*>(ctx->managed);
if (mt->deleter != nullptr) {
mt->deleter(mt);
}
} else {
auto* mt = static_cast<dlpack_abi::DLManagedTensor*>(ctx->managed);
if (mt->deleter != nullptr) {
mt->deleter(mt);
}
}
delete ctx;
}
}
} // namespace
Array TryFromDlpackZerocopy(nb::handle obj) {
if (!nb::hasattr(obj, "__dlpack__")) {
return Array();
}
nb::object cap;
try {
cap = obj.attr("__dlpack__")();
} catch (const std::exception&) {
return Array();
}
if (!PyCapsule_CheckExact(cap.ptr())) {
return Array();
}
const char* name = PyCapsule_GetName(cap.ptr());
if (name == nullptr) {
PyErr_Clear();
return Array();
}
const bool versioned = std::strcmp(name, "dltensor_versioned") == 0;
if (!versioned && std::strcmp(name, "dltensor") != 0) {
return Array();
}
void* ptr = PyCapsule_GetPointer(cap.ptr(), name);
if (ptr == nullptr) {
PyErr_Clear();
return Array();
}
dlpack_abi::DLTensor* t = nullptr;
std::function<void()> deleter;
if (versioned) {
auto* mt = static_cast<dlpack_abi::DLManagedTensorVersioned*>(ptr);
t = &mt->dl_tensor;
deleter = [mt]() {
if (mt->deleter != nullptr) {
mt->deleter(mt);
}
};
} else {
auto* mt = static_cast<dlpack_abi::DLManagedTensor*>(ptr);
t = &mt->dl_tensor;
deleter = [mt]() {
if (mt->deleter != nullptr) {
mt->deleter(mt);
}
};
}
if (t == nullptr || t->data == nullptr) {
return Array();
}
DType dt;
try {
nb::dlpack::dtype d;
d.lanes = t->dtype.lanes;
d.code = t->dtype.code;
d.bits = t->dtype.bits;
dt = FromDlpackDtype(d);
} catch (const std::exception&) {
return Array();
}
std::vector<int64_t> dims(static_cast<size_t>(t->ndim));
for (int32_t i = 0; i < t->ndim; ++i) {
dims[static_cast<size_t>(i)] = t->shape[i];
}
Shape shape(dims);
const bool gpu_device = (t->device.device_type == kDLCUDA_DEV || t->device.device_type == kDLROCM_DEV || t->device.device_type == kDLOneAPI_DEV);
if (gpu_device) {
// Zero-copy device import via a non-owned PJRT view; requires a dense layout.
if (!PlatformIsGpu() || !DeviceImportSupported() || t->byte_offset != 0) {
return Array();
}
if (t->strides != nullptr) {
const std::vector<int64_t> expect = shape.ComputeContiguousStrides();
for (int32_t i = 0; i < t->ndim; ++i) {
if (t->strides[i] != expect[static_cast<size_t>(i)]) {
return Array();
}
}
}
PyCapsule_SetName(cap.ptr(), "used_dltensor");
auto* ctx = new DlpackImportCtx{versioned, ptr};
const auto& devs = HardwareManager::Instance().Devices();
PJRT_Device* dev = nullptr;
if (!devs.empty()) {
const int did = t->device.device_id;
dev = (did >= 0 && did < static_cast<int>(devs.size())) ? devs[static_cast<size_t>(did)].handle : devs.front().handle;
}
std::shared_ptr<DeviceBuffer> db;
try {
db = DeviceBuffer::CreateView(PjrtApi(), HardwareManager::Instance().Client(), dev, t->data, dims, ToPjrtBufferType(dt), 0, &ImportViewDeleter, ctx);
} catch (const std::exception&) {
delete ctx;
return Array();
}
Array arr = Array::Parameter(shape, dt);
arr.BindDeviceData(db);
return arr;
}
// Host (CPU) memory is wrapped zero-copy by a PinnedBuffer.
if (t->device.device_type != dlpack_abi::kDLCPU) {
return Array();
}
// Commit: rename the capsule so the producer no longer frees the tensor.
PyCapsule_SetName(cap.ptr(), "used_dltensor");
const size_t elem = SizeOf(dt);
const int64_t offset_elems = elem != 0 ? static_cast<int64_t>(t->byte_offset / elem) : 0;
std::vector<int64_t> strides;
if (t->strides != nullptr) {
strides.reserve(static_cast<size_t>(t->ndim));
for (int32_t i = 0; i < t->ndim; ++i) {
strides.push_back(t->strides[i]);
}
} else {
strides = shape.ComputeContiguousStrides();
}
const size_t bytes = static_cast<size_t>(shape.NumElements()) * elem;
auto buf = PinnedBuffer::Wrap(t->data, bytes, std::move(deleter));
Array arr = Array::Parameter(shape, dt);
arr.BindData(buf);
return arr.AsStrided(shape, strides, offset_elems);
}
inline nb::object ElementToPy(const void* p, DType dt) {
switch (dt) {
case DType::kBool:
return nb::cast(*reinterpret_cast<const uint8_t*>(p) != 0);
case DType::kInt8:
return nb::cast(static_cast<int>(*reinterpret_cast<const int8_t*>(p)));
case DType::kUInt8:
return nb::cast(static_cast<int>(*reinterpret_cast<const uint8_t*>(p)));
case DType::kInt16:
return nb::cast(static_cast<int>(*reinterpret_cast<const int16_t*>(p)));
case DType::kUInt16:
return nb::cast(static_cast<int>(*reinterpret_cast<const uint16_t*>(p)));
case DType::kInt32:
return nb::cast(*reinterpret_cast<const int32_t*>(p));
case DType::kUInt32:
return nb::cast(*reinterpret_cast<const uint32_t*>(p));
case DType::kInt64:
return nb::cast(*reinterpret_cast<const int64_t*>(p));
case DType::kUInt64:
return nb::cast(*reinterpret_cast<const uint64_t*>(p));
case DType::kFloat16:
return nb::cast(static_cast<double>(*reinterpret_cast<const tenzr::float16*>(p)));
case DType::kBFloat16:
return nb::cast(static_cast<double>(*reinterpret_cast<const tenzr::bfloat16*>(p)));
case DType::kFloat32:
return nb::cast(static_cast<double>(*reinterpret_cast<const float*>(p)));
case DType::kFloat64:
return nb::cast(*reinterpret_cast<const double*>(p));
case DType::kFloat8E4M3FN:
return nb::cast(static_cast<double>(*reinterpret_cast<const tenzr::float8_e4m3fn*>(p)));
case DType::kFloat8E5M2:
return nb::cast(static_cast<double>(*reinterpret_cast<const tenzr::float8_e5m2*>(p)));
case DType::kFloat8E3M4:
return nb::cast(static_cast<double>(*reinterpret_cast<const tenzr::float8_e3m4*>(p)));
case DType::kFloat8E4M3:
return nb::cast(static_cast<double>(*reinterpret_cast<const tenzr::float8_e4m3*>(p)));
case DType::kFloat8E4M3FNUZ:
return nb::cast(static_cast<double>(*reinterpret_cast<const tenzr::float8_e4m3fnuz*>(p)));
case DType::kFloat8E4M3B11FNUZ:
return nb::cast(static_cast<double>(*reinterpret_cast<const tenzr::float8_e4m3b11fnuz*>(p)));
case DType::kFloat8E5M2FNUZ:
return nb::cast(static_cast<double>(*reinterpret_cast<const tenzr::float8_e5m2fnuz*>(p)));
case DType::kFloat8E8M0FNU:
return nb::cast(static_cast<double>(*reinterpret_cast<const tenzr::float8_e8m0fnu*>(p)));
case DType::kFloat6E2M3FN:
return nb::cast(static_cast<double>(*reinterpret_cast<const tenzr::float6_e2m3fn*>(p)));
case DType::kFloat6E3M2FN:
return nb::cast(static_cast<double>(*reinterpret_cast<const tenzr::float6_e3m2fn*>(p)));
case DType::kFloat4E2M1FN:
return nb::cast(static_cast<double>(*reinterpret_cast<const tenzr::float4_e2m1fn*>(p)));
case DType::kInt1:
return nb::cast(static_cast<int64_t>(*reinterpret_cast<const tenzr::int1*>(p)));
case DType::kUInt1:
return nb::cast(static_cast<int64_t>(*reinterpret_cast<const tenzr::uint1*>(p)));
case DType::kInt2:
return nb::cast(static_cast<int64_t>(*reinterpret_cast<const tenzr::int2*>(p)));
case DType::kUInt2:
return nb::cast(static_cast<int64_t>(*reinterpret_cast<const tenzr::uint2*>(p)));
case DType::kInt4:
return nb::cast(static_cast<int64_t>(*reinterpret_cast<const tenzr::int4*>(p)));
case DType::kUInt4:
return nb::cast(static_cast<int64_t>(*reinterpret_cast<const tenzr::uint4*>(p)));
default:
TENZR_THROW("Unsupported dtype in ElementToPy");
}
}
nb::object BuildNestedList(const uint8_t* base, const StridedLayout& lay, const Array& a, size_t dim, int64_t offset, size_t elem) {
if (dim == lay.shape().Rank()) {
return ElementToPy(base + static_cast<size_t>(offset) * elem, a.Dtype());
}
nb::list out;
const int64_t stride = lay.Stride(dim);
const int64_t n = lay.shape().Dim(dim);
for (int64_t i = 0; i < n; ++i) {
out.append(BuildNestedList(base, lay, a, dim + 1, offset + i * stride, elem));
}
return out;
}
nb::object ArrayTolist(const Array& a_orig) {
Array a = a_orig;
if (!a.HasData() && !a.IsSymbolic()) {
a = a.Eval();
}
TENZR_CHECK(a.HasData(), "cannot convert a symbolic array to a Python list");
a.SyncHost();
const uint8_t* base = static_cast<const uint8_t*>(a.Data()->Data());
return BuildNestedList(base, a.Layout(), a, 0, a.Layout().Offset(), SizeOf(a.Dtype()));
}
nb::object ArrayScalar(const Array& a) {
TENZR_CHECK(a.shape().NumElements() == 1, "item() requires a scalar array with exactly 1 element");
switch (a.Dtype()) {
case DType::kBool:
return nb::cast(static_cast<bool>(a.ToVector<bool>()[0]));
case DType::kInt8:
return nb::cast(a.ToVector<int8_t>()[0]);
case DType::kUInt8:
return nb::cast(a.ToVector<uint8_t>()[0]);
case DType::kInt16:
return nb::cast(a.ToVector<int16_t>()[0]);
case DType::kUInt16:
return nb::cast(a.ToVector<uint16_t>()[0]);
case DType::kInt32:
return nb::cast(a.ToVector<int32_t>()[0]);
case DType::kUInt32:
return nb::cast(a.ToVector<uint32_t>()[0]);
case DType::kInt64:
return nb::cast(a.ToVector<int64_t>()[0]);
case DType::kUInt64:
return nb::cast(a.ToVector<uint64_t>()[0]);
case DType::kFloat16:
return nb::cast(static_cast<double>(a.ToVector<tenzr::float16>()[0]));
case DType::kBFloat16:
return nb::cast(static_cast<double>(a.ToVector<tenzr::bfloat16>()[0]));
case DType::kFloat32:
return nb::cast(static_cast<double>(a.ToVector<float>()[0]));
case DType::kFloat64:
return nb::cast(a.ToVector<double>()[0]);
case DType::kFloat8E4M3FN:
return nb::cast(static_cast<double>(a.ToVector<tenzr::float8_e4m3fn>()[0]));
case DType::kFloat8E5M2:
return nb::cast(static_cast<double>(a.ToVector<tenzr::float8_e5m2>()[0]));
case DType::kFloat8E3M4:
return nb::cast(static_cast<double>(a.ToVector<tenzr::float8_e3m4>()[0]));
case DType::kFloat8E4M3:
return nb::cast(static_cast<double>(a.ToVector<tenzr::float8_e4m3>()[0]));
case DType::kFloat8E4M3FNUZ:
return nb::cast(static_cast<double>(a.ToVector<tenzr::float8_e4m3fnuz>()[0]));
case DType::kFloat8E4M3B11FNUZ:
return nb::cast(static_cast<double>(a.ToVector<tenzr::float8_e4m3b11fnuz>()[0]));
case DType::kFloat8E5M2FNUZ:
return nb::cast(static_cast<double>(a.ToVector<tenzr::float8_e5m2fnuz>()[0]));
case DType::kFloat8E8M0FNU:
return nb::cast(static_cast<double>(a.ToVector<tenzr::float8_e8m0fnu>()[0]));
case DType::kFloat6E2M3FN:
return nb::cast(static_cast<double>(a.ToVector<tenzr::float6_e2m3fn>()[0]));
case DType::kFloat6E3M2FN:
return nb::cast(static_cast<double>(a.ToVector<tenzr::float6_e3m2fn>()[0]));
case DType::kFloat4E2M1FN:
return nb::cast(static_cast<double>(a.ToVector<tenzr::float4_e2m1fn>()[0]));
case DType::kInt1:
return nb::cast(static_cast<int64_t>(a.ToVector<tenzr::int1>()[0]));
case DType::kUInt1:
return nb::cast(static_cast<int64_t>(a.ToVector<tenzr::uint1>()[0]));
case DType::kInt2:
return nb::cast(static_cast<int64_t>(a.ToVector<tenzr::int2>()[0]));
case DType::kUInt2:
return nb::cast(static_cast<int64_t>(a.ToVector<tenzr::uint2>()[0]));
case DType::kInt4:
return nb::cast(static_cast<int64_t>(a.ToVector<tenzr::int4>()[0]));
case DType::kUInt4:
return nb::cast(static_cast<int64_t>(a.ToVector<tenzr::uint4>()[0]));
default:
TENZR_THROW("Unsupported dtype in item()");
}
}
std::string ArrayRepr(const Array& a) {
std::ostringstream ss;
ss << "Array(";
if (a.IsSymbolic()) {
ss << "symbolic: var_" << a.SymbolicVarId().id << ", ";
} else {
ss << nb::cast<std::string>(nb::str(ArrayTolist(a))) << ", ";
}
ss << "shape=" << a.shape().ToString() << ", dtype=" << ToString(a.Dtype()) << ")";
return ss.str();
}
void BindDTypeEnum(nb::module_& m) {
// Enum DType
nb::enum_<DType>(m, "dtype")
.value("bool_", DType::kBool)
.value("float32", DType::kFloat32)
.value("float64", DType::kFloat64)
.value("int32", DType::kInt32)
.value("int64", DType::kInt64)
.value("float16", DType::kFloat16)
.value("bfloat16", DType::kBFloat16)
.value("float8_e4m3fn", DType::kFloat8E4M3FN)
.value("float8_e5m2", DType::kFloat8E5M2)
.value("uint8", DType::kUInt8)
.value("int8", DType::kInt8)
.value("int16", DType::kInt16)
.value("uint16", DType::kUInt16)
.value("uint32", DType::kUInt32)
.value("uint64", DType::kUInt64)
.value("float8_e3m4", DType::kFloat8E3M4)
.value("float8_e4m3", DType::kFloat8E4M3)
.value("float8_e4m3fnuz", DType::kFloat8E4M3FNUZ)
.value("float8_e4m3b11fnuz", DType::kFloat8E4M3B11FNUZ)
.value("float8_e5m2fnuz", DType::kFloat8E5M2FNUZ)
.value("float8_e8m0fnu", DType::kFloat8E8M0FNU)
.value("float6_e2m3fn", DType::kFloat6E2M3FN)
.value("float6_e3m2fn", DType::kFloat6E3M2FN)
.value("float4_e2m1fn", DType::kFloat4E2M1FN)
.value("int1", DType::kInt1)
.value("uint1", DType::kUInt1)
.value("int2", DType::kInt2)
.value("uint2", DType::kUInt2)
.value("int4", DType::kInt4)
.value("uint4", DType::kUInt4)
.export_values();
}
void BindShape(nb::module_& m) {
// Shape class
nb::class_<Shape>(m, "Shape")
.def(nb::init<std::vector<int64_t>>())
.def_prop_ro("rank", &Shape::Rank)
.def_prop_ro("num_elements", &Shape::NumElements)
.def("__len__", &Shape::Rank)
.def("__getitem__",
[](const Shape& s, size_t i) {
if (i >= s.Rank()) {
throw nb::index_error();
}
return s.Dim(i);
})
.def("__repr__", [](const Shape& s) {
std::ostringstream ss;
ss << "(";
for (size_t i = 0; i < s.Rank(); ++i) {
ss << s.Dim(i) << (i + 1 < s.Rank() || s.Rank() == 1 ? "," : "");
if (i + 1 < s.Rank()) {
ss << " ";
}
}
ss << ")";
return ss.str();
});
}
void BindArrayModuleFunctions(nb::module_& m) {
// Module-level array creation functions
m.def("array", [](nb::handle obj, std::optional<DType> dtype) -> Array { return PyToArray(obj, dtype); }, "object"_a, "dtype"_a = nb::none());
m.def(
"from_dlpack",
[](nb::handle obj) -> Array {
if (nb::isinstance<Array>(obj)) {
return nb::cast<Array>(obj);
}
if (nb::hasattr(obj, "__dlpack__")) {
Array zero = TryFromDlpackZerocopy(obj);
if (zero.IsDefined()) {
return zero;
}
nb::object np = nb::module_::import_("numpy");
return PyToArray(np.attr("from_dlpack")(obj));
}
if (nb::hasattr(obj, "__array__")) {
nb::object np = nb::module_::import_("numpy");
return PyToArray(np.attr("asarray")(obj));
}
throw nb::type_error("from_dlpack() expects an object implementing __dlpack__ or __array__");
},
"object"_a);
m.def("zeros", [](std::vector<int64_t> shape, DType dtype) -> Array { return Array::Zeros(Shape(shape), dtype); }, "shape"_a, "dtype"_a = DType::kFloat32);
m.def("ones", [](std::vector<int64_t> shape, DType dtype) -> Array { return Array::Ones(Shape(shape), dtype); }, "shape"_a, "dtype"_a = DType::kFloat32);
m.def(
"full", [](std::vector<int64_t> shape, double val, DType dtype) -> Array { return Array::Full(Shape(shape), val, dtype); }, "shape"_a, "fill_value"_a,
"dtype"_a = DType::kFloat32);
m.def(
"arange",
[](double start, std::optional<double> stop, double step, DType dtype) -> Array {
double actual_start = stop.has_value() ? start : 0.0;
double actual_stop = stop.has_value() ? *stop : start;
return tenzr::numpy::Arange(actual_start, actual_stop, step, dtype);
},
"start"_a, "stop"_a = nb::none(), "step"_a = 1.0, "dtype"_a = DType::kFloat32);
m.def(
"linspace", [](double start, double stop, int64_t num, DType dtype) -> Array { return tenzr::numpy::Linspace(start, stop, num, dtype); }, "start"_a,
"stop"_a, "num"_a = 50, "dtype"_a = DType::kFloat32);
m.def(
"eye", [](int64_t N, std::optional<int64_t> M, DType dtype) -> Array { return tenzr::numpy::Eye(N, M.value_or(N), 0, dtype); }, "N"_a,
"M"_a = nb::none(), "dtype"_a = DType::kFloat32);
m.def(
"concatenate", [](const std::vector<Array>& arrays, int64_t axis) -> Array { return tenzr::numpy::Concatenate(arrays, axis); }, "arrays"_a,
"axis"_a = 0);
m.def("stack", [](const std::vector<Array>& arrays, int64_t axis) -> Array { return tenzr::numpy::Stack(arrays, axis); }, "arrays"_a, "axis"_a = 0);
m.def(
"where", [](const Array& condition, const Array& x, const Array& y) -> Array { return tenzr::numpy::Where(condition, x, y); }, "condition"_a, "x"_a,
"y"_a);
// Diagnostic: true when a trace is active (jit/vmap symbolic tracing or grad tracing).
m.def("_is_tracing", []() { return txpr::TracerContext::IsActive() || txpr::IsGradTracingActive(); });
// Diagnostic: PJRT DLPack capability report (platform, api, pointer/refcount/view support).
m.def("_pjrt_dlpack_caps", []() {
const PJRT_Api* api = PjrtApi();
return nb::make_tuple(LowerPlatform(), api != nullptr, api != nullptr && api->PJRT_Buffer_UnsafePointer != nullptr,
api != nullptr && api->PJRT_Buffer_IncreaseExternalReferenceCount != nullptr,
api != nullptr && api->PJRT_Buffer_DecreaseExternalReferenceCount != nullptr,
api != nullptr && api->PJRT_Client_CreateViewOfDeviceBuffer != nullptr);
});
}
void BindArray(nb::module_& m) {
BindDTypeEnum(m);
BindShape(m);
// Array class
nb::class_<Array> arr(m, "Array");
arr.def_prop_ro("shape",
[](const Array& a) {
std::vector<int64_t> dims(a.shape().Rank());
for (size_t i = 0; i < a.shape().Rank(); ++i) {
dims[i] = a.shape().Dim(i);
}
return nb::cast(dims);
})
.def_prop_ro("dtype", [](const Array& a) { return a.Dtype(); })
.def_prop_ro("ndim", [](const Array& a) { return a.shape().Rank(); })
.def_prop_ro("size", [](const Array& a) { return a.shape().NumElements(); })
.def_prop_ro("is_symbolic", &Array::IsSymbolic)
.def_prop_ro("has_device_data", &Array::HasDeviceData)
.def_prop_ro("has_host_data", &Array::HasData)
.def_prop_ro("nbytes", [](const Array& a) { return a.shape().NumElements() * static_cast<int64_t>(SizeOf(a.Dtype())); })
.def_prop_ro("itemsize", [](const Array& a) { return static_cast<int64_t>(SizeOf(a.Dtype())); })
.def_prop_ro("T",
[](const Array& a) -> Array {
int64_t r = static_cast<int64_t>(a.shape().Rank());
if (r < 2) {
return a;
}
std::vector<int64_t> perm(static_cast<size_t>(r));
for (int64_t i = 0; i < r; ++i) {
perm[static_cast<size_t>(i)] = r - 1 - i;
}
return a.Permute(perm);
})
.def_prop_ro("mT", [](const Array& a) -> Array { return a.shape().Rank() >= 2 ? a.Transpose() : a; })
.def_prop_ro("device",
[](const Array&) {
const std::string& p = HardwareManager::Instance().PlatformName();
return p.empty() ? std::string("cpu") : p;
})
.def_prop_ro("devices",
[](const Array&) {
const std::string& p = HardwareManager::Instance().PlatformName();
return std::vector<std::string>{p.empty() ? std::string("cpu") : p};
})
.def("__len__",
[](const Array& a) -> int64_t {
if (a.shape().Rank() == 0) {
throw nb::type_error("len() of unsized object");
}
return a.shape().Dim(0);
})
.def("__iter__",
[](const Array& a) -> nb::object {
if (a.shape().Rank() == 0) {
throw nb::type_error("iteration over a 0-d array");
}
nb::list rows;
for (int64_t i = 0; i < a.shape().Dim(0); ++i) {
rows.append(ArrayGetitem(a, nb::int_(i)));
}
return nb::iter(rows);
})
.def("__hash__", [](const Array&) -> int64_t { throw nb::type_error("unhashable type: 'tenzr.Array'"); })
.def("copy", [](const Array& a) -> Array { return a.Clone(); })
.def("eval", [](const Array& a) -> Array { return a.Eval(); })
.def(
"__array__",
[](const Array& a, nb::handle dtype, nb::handle copy) {
nb::object nd = ArrayToNdarray(a, nb::handle());
if (!dtype.is_none()) {
nd = nd.attr("astype")(dtype);
}
if (!copy.is_none() && nb::cast<bool>(copy)) {
nd = nd.attr("copy")();
}
return nd;
},
"dtype"_a = nb::none(), "copy"_a = nb::none())
.def_prop_ro("__array_interface__",
[](nb::handle self) {
const Array& a = nb::cast<const Array&>(self);
// Build the interface dict directly: numpy 2.5's
// ndarray.__array_interface__ getter (array_protocol_descr_get)
// segfaults on nanobind-created arrays.
Array export_arr = a;
if (!export_arr.HasData() && !export_arr.IsSymbolic()) {
export_arr = export_arr.Eval();
}
TENZR_CHECK(export_arr.HasData(), "Cannot export an unevaluated symbolic array to Python");
export_arr.SyncHost();
const StridedLayout& lay = export_arr.Layout();
const size_t rank = lay.shape().Rank();
const size_t elem = SizeOf(export_arr.Dtype());
const uintptr_t base = reinterpret_cast<uintptr_t>(export_arr.Data()->Data()) + static_cast<uintptr_t>(lay.Offset()) * elem;
nb::list shape;
nb::list strides;
for (size_t i = 0; i < rank; ++i) {
shape.append(lay.shape().Dim(i));
strides.append(lay.Stride(i) * static_cast<int64_t>(elem));
}
nb::object nd = ArrayToNdarray(a, self);
nb::dict iface;
iface["data"] = nb::make_tuple(base, false);
iface["shape"] = nb::steal<nb::object>(PyList_AsTuple(shape.ptr()));
iface["strides"] = (rank == 0) ? nb::none() : nb::steal<nb::object>(PyList_AsTuple(strides.ptr()));
iface["typestr"] = nd.attr("dtype").attr("str");
iface["version"] = 3;
return iface;
})
.def(
"__array_namespace__", [](nb::handle /*self*/, nb::handle /*api_version*/) { return nb::module_::import_("tenzr.numpy"); },
"api_version"_a = nb::none())
.def(
"__buffer__",
[](nb::handle self, int /*flags*/) {
const Array& a = nb::cast<const Array&>(self);
return nb::memoryview(ArrayToNdarray(a, self));
},
"flags"_a = 0)
.def("__repr__", [](const Array& a) { return ArrayRepr(a); })
.def("__str__",
[](const Array& a) {
if (!a.HasData() && a.IsSymbolic()) {
std::ostringstream ss;
ss << "Array(shape=(";
for (size_t i = 0; i < a.shape().Rank(); ++i) {
ss << a.shape().Dim(i) << (i + 1 < a.shape().Rank() || a.shape().Rank() == 1 ? "," : "");
}
ss << "), dtype=" << ToString(a.Dtype()) << ")";
return ss.str();
}
return nb::cast<std::string>(nb::str(ArrayTolist(a)));
})
// Buffer protocol & DLPack
.def("numpy",
[](nb::handle self) {
const Array& a = nb::cast<const Array&>(self);
return ArrayToNdarray(a, self);
})
.def("__dlpack__",
[](nb::handle self, nb::args args, nb::kwargs kwargs) -> nb::object {
const Array& a = nb::cast<const Array&>(self);
if (a.HasDeviceData() && PlatformIsGpu() && DeviceExportSupported()) {
return MakeDeviceDlpack(a);
}
auto nd = ArrayToNdarray(a, self);
PyObject* fn = nd.attr("__dlpack__").ptr();
PyObject* res = PyObject_Call(fn, args.ptr(), kwargs.ptr());
if (!res) {
throw nb::python_error();
}
return nb::steal<nb::object>(res);
})
.def("__dlpack_device__",
[](const Array& a) {
if (a.HasDeviceData() && PlatformIsGpu() && DeviceExportSupported()) {
int id = a.DeviceData() ? a.DeviceData()->DeviceId() : 0;
return nb::make_tuple(DlDeviceTypeForPlatform(), id);
}
return nb::make_tuple(1, 0); // CPU device (kDLCPU = 1, id = 0)
})
.def("tolist", [](const Array& a) { return ArrayTolist(a); })
.def("item", [](const Array& a) -> nb::object { return ArrayScalar(a); })
// Python scalar protocols (NumPy/JAX semantics): truthiness and
// scalar conversion implicitly materialize, while symbolic arrays raise.
.def("__bool__",
[](const Array& a) -> bool {
const int64_t n = a.shape().NumElements();
if (n == 0) {
throw nb::value_error("The truth value of an empty array is ambiguous. Use a.size > 0 to test for emptiness.");
}
if (n > 1) {
throw nb::value_error("The truth value of an array with more than one element is ambiguous. Use a.any() or a.all()");
}
return PyObject_IsTrue(ArrayScalar(a).ptr()) != 0;
})
.def("__float__",
[](const Array& a) -> double {
if (a.shape().Rank() != 0) {
throw nb::type_error(("Only scalar arrays can be converted to Python scalars; got ndim=" + std::to_string(a.shape().Rank())).c_str());
}
double v = PyFloat_AsDouble(ArrayScalar(a).ptr());
if (PyErr_Occurred()) {
throw nb::python_error();
}
return v;
})
.def("__int__",
[](const Array& a) -> int64_t {
if (a.shape().Rank() != 0) {
throw nb::type_error(("Only scalar arrays can be converted to Python scalars; got ndim=" + std::to_string(a.shape().Rank())).c_str());
}
PyObject* v = PyNumber_Long(ArrayScalar(a).ptr());
if (v == nullptr) {
throw nb::python_error();
}
return nb::cast<int64_t>(nb::steal<nb::object>(v));
})
.def("__complex__",
[](const Array& a) -> nb::object {
if (a.shape().Rank() != 0) {
throw nb::type_error(("Only scalar arrays can be converted to Python scalars; got ndim=" + std::to_string(a.shape().Rank())).c_str());
}
Py_complex c = PyComplex_AsCComplex(ArrayScalar(a).ptr());
if (PyErr_Occurred()) {
throw nb::python_error();
}
return nb::steal<nb::object>(PyComplex_FromDoubles(c.real, c.imag));
})
.def("__index__",
[](const Array& a) -> int64_t {
if (a.shape().Rank() != 0 || !IsIntegral(a.Dtype())) {
throw nb::type_error("Only integer scalar arrays can be converted to a scalar index");
}
return nb::cast<int64_t>(ArrayScalar(a));
})
// Binary Arithmetic (JAX-style weak-type promotion; D6)
.def("__add__", [](const Array& a, nb::handle other) { return BinaryArith(a, other, false, [](const Array& x, const Array& y) { return x + y; }); })
.def("__radd__", [](const Array& a, nb::handle other) { return BinaryArith(a, other, false, [](const Array& x, const Array& y) { return y + x; }); })
.def("__sub__", [](const Array& a, nb::handle other) { return BinaryArith(a, other, false, [](const Array& x, const Array& y) { return x - y; }); })
.def("__rsub__", [](const Array& a, nb::handle other) { return BinaryArith(a, other, false, [](const Array& x, const Array& y) { return y - x; }); })
.def("__mul__", [](const Array& a, nb::handle other) { return BinaryArith(a, other, false, [](const Array& x, const Array& y) { return x * y; }); })
.def("__rmul__", [](const Array& a, nb::handle other) { return BinaryArith(a, other, false, [](const Array& x, const Array& y) { return y * x; }); })
.def("__truediv__", [](const Array& a, nb::handle other) { return BinaryArith(a, other, true, [](const Array& x, const Array& y) { return x / y; }); })
.def("__rtruediv__", [](const Array& a, nb::handle other) { return BinaryArith(a, other, true, [](const Array& x, const Array& y) { return y / x; }); })
.def("__neg__", [](const Array& a) { return -a; })
.def("__matmul__", [](const Array& a, nb::handle other) { return tenzr::numpy::Matmul(a, PyToArray(other, a.Dtype())); })
// Comparisons (promoted; D6)
.def("__eq__", [](const Array& a, nb::handle other) { return BinaryCmp(a, other, [](const Array& x, const Array& y) { return x == y; }); })
.def("__ne__", [](const Array& a, nb::handle other) { return BinaryCmp(a, other, [](const Array& x, const Array& y) { return x != y; }); })
.def("__lt__", [](const Array& a, nb::handle other) { return BinaryCmp(a, other, [](const Array& x, const Array& y) { return x < y; }); })
.def("__le__", [](const Array& a, nb::handle other) { return BinaryCmp(a, other, [](const Array& x, const Array& y) { return x <= y; }); })
.def("__gt__", [](const Array& a, nb::handle other) { return BinaryCmp(a, other, [](const Array& x, const Array& y) { return x > y; }); })
.def("__ge__", [](const Array& a, nb::handle other) { return BinaryCmp(a, other, [](const Array& x, const Array& y) { return x >= y; }); })
// Slicing & Indexing
.def("__getitem__", [](const Array& a, nb::handle key) -> Array { return ArrayGetitem(a, key); })
// Immutability guard: Array is a functional expression node; in-place assignment
// would invalidate AST hashing / JIT caching and the autograd tape.
.def("__setitem__",
[](Array&, nb::handle, nb::handle) {
throw nb::type_error(
"tenzr.Array is immutable and does not support item assignment; "
"build a modified copy with functional updates (e.g. tenzr.where / plax.select)");
})
// JAX-style functional indexed updates: a.at[idx].get/set/add/multiply/min/max.
.def_prop_ro("at", [](nb::handle self) { return nb::module_::import_("tenzr._indexing").attr("_IndexUpdateHelper")(self); })
// Transforms & Views
.def("reshape",
[](const Array& a, nb::args args) -> Array {
std::vector<int64_t> new_dims;
if (args.size() == 1 && (nb::isinstance<nb::tuple>(args[0]) || nb::isinstance<nb::list>(args[0]))) {
for (auto item : nb::cast<nb::sequence>(args[0])) {
new_dims.push_back(nb::cast<int64_t>(item));
}
} else {
for (size_t i = 0; i < args.size(); ++i) {
new_dims.push_back(nb::cast<int64_t>(args[i]));
}
}
// NumPy/JAX-style single -1 inference.
int64_t neg_index = -1;
int64_t known_product = 1;
for (size_t i = 0; i < new_dims.size(); ++i) {
if (new_dims[i] == -1) {
TENZR_CHECK(neg_index == -1, "reshape accepts at most one inferred (-1) dimension");
neg_index = static_cast<int64_t>(i);
} else {
TENZR_CHECK(new_dims[i] >= 0, "reshape dimensions must be >= 0 or exactly one -1");
known_product *= new_dims[i];
}
}
if (neg_index >= 0) {
const int64_t total = a.shape().NumElements();
TENZR_CHECK(known_product > 0 && total % known_product == 0, "cannot infer -1 dimension: size is not divisible");
new_dims[static_cast<size_t>(neg_index)] = total / known_product;
}
return a.Reshape(Shape(new_dims));
})
.def(
"transpose",
[](const Array& a, std::optional<std::vector<int64_t>> perm) -> Array {
if (perm.has_value()) {
return plax::Transpose(a, *perm);
}
return a.Transpose();
},
"axes"_a = nb::none())
.def(
"squeeze",
[](const Array& a, std::optional<int64_t> axis) -> Array {
if (axis.has_value()) {
return a.Squeeze(static_cast<size_t>(*axis));
}
return a.Squeeze();
},
"axis"_a = nb::none())
.def(
"expand_dims", [](const Array& a, int64_t axis) -> Array { return a.Unsqueeze(axis); }, "axis"_a)
.def(
"unsqueeze", [](const Array& a, int64_t axis) -> Array { return a.Unsqueeze(axis); }, "axis"_a)
.def(
"astype", [](const Array& a, DType new_dt) -> Array { return plax::ConvertElementType(a, new_dt); }, "dtype"_a)
.def("flatten", [](const Array& a) -> Array { return a.Reshape(Shape({a.shape().NumElements()})); })
// Reductions
.def(
"sum", [](const Array& a, nb::handle axis, bool keepdims) -> Array { return plax::ReduceSum(a, ParseAxes(axis), keepdims); }, "axis"_a = nb::none(),
"keepdims"_a = false)
.def(
"mean", [](const Array& a, nb::handle axis, bool keepdims) -> Array { return tenzr::numpy::Mean(a, ParseAxes(axis), keepdims); },
"axis"_a = nb::none(), "keepdims"_a = false)
.def(
"max", [](const Array& a, nb::handle axis, bool keepdims) -> Array { return plax::ReduceMax(a, ParseAxes(axis), keepdims); }, "axis"_a = nb::none(),
"keepdims"_a = false)
.def(
"min", [](const Array& a, nb::handle axis, bool keepdims) -> Array { return plax::ReduceMin(a, ParseAxes(axis), keepdims); }, "axis"_a = nb::none(),
"keepdims"_a = false)
.def(
"prod", [](const Array& a, nb::handle axis, bool keepdims) -> Array { return a.Prod(ParseAxes(axis), keepdims); }, "axis"_a = nb::none(),
"keepdims"_a = false)
.def(
"var", [](const Array& a, nb::handle axis, bool keepdims) -> Array { return tenzr::numpy::Var(a, ParseAxes(axis), keepdims); },
"axis"_a = nb::none(), "keepdims"_a = false)
.def(
"std", [](const Array& a, nb::handle axis, bool keepdims) -> Array { return tenzr::numpy::StdDev(a, ParseAxes(axis), keepdims); },
"axis"_a = nb::none(), "keepdims"_a = false)
.def("ravel", [](const Array& a) -> Array { return a.Reshape(Shape({a.shape().NumElements()})); })
.def(
"broadcast_to", [](const Array& a, std::vector<int64_t> shape) -> Array { return a.BroadcastTo(Shape(shape)); }, "shape"_a)
.def(
"take", [](const Array& a, const Array& indices, int64_t axis) -> Array { return tenzr::numpy::Take(a, indices, axis); }, "indices"_a, "axis"_a = 0)
.def(
"power", [](const Array& a, nb::handle other) -> Array { return a.Pow(PyToArray(other, a.Dtype())); }, "other"_a)
.def(
"maximum", [](const Array& a, nb::handle other) -> Array { return a.Maximum(PyToArray(other, a.Dtype())); }, "other"_a)
.def(
"minimum", [](const Array& a, nb::handle other) -> Array { return a.Minimum(PyToArray(other, a.Dtype())); }, "other"_a)
.def(
"remainder", [](const Array& a, nb::handle other) -> Array { return a.Remainder(PyToArray(other, a.Dtype())); }, "other"_a)
.def(
"atan2", [](const Array& a, nb::handle other) -> Array { return a.Atan2(PyToArray(other, a.Dtype())); }, "other"_a)
.def("tan", [](const Array& a) { return a.Tan(); })
.def("ceil", [](const Array& a) { return a.Ceil(); })
.def("floor", [](const Array& a) { return a.Floor(); })
.def("round", [](const Array& a) { return a.Round(); })
.def("sign", [](const Array& a) { return a.Sign(); })
.def("cbrt", [](const Array& a) { return a.Cbrt(); })
.def("expm1", [](const Array& a) { return a.Expm1(); })
.def("log1p", [](const Array& a) { return a.Log1p(); })
.def("is_finite", [](const Array& a) { return a.IsFinite(); })
.def("logical_not", [](const Array& a) { return a.LogicalNot(); })
.def(
"bitwise_and", [](const Array& a, nb::handle o) -> Array { return a.BitwiseAnd(PyToArray(o, a.Dtype())); }, "other"_a)
.def(
"bitwise_or", [](const Array& a, nb::handle o) -> Array { return a.BitwiseOr(PyToArray(o, a.Dtype())); }, "other"_a)
.def(
"bitwise_xor", [](const Array& a, nb::handle o) -> Array { return a.BitwiseXor(PyToArray(o, a.Dtype())); }, "other"_a)
.def(
"shift_left", [](const Array& a, nb::handle o) -> Array { return a.ShiftLeft(PyToArray(o, a.Dtype())); }, "other"_a)
.def("shift_right", [](const Array& a, nb::handle o) -> Array { return a.ShiftRight(PyToArray(o, a.Dtype())); }, "other"_a);
// Prefer Tenzr Array over numpy ndarray when the two are combined, matching JAX.
arr.attr("__array_priority__") = 1000;
BindArrayModuleFunctions(m);
}
} // namespace tenzr::python