src/python/bindings.h¶

// 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.

#ifndef TENZR_SRC_PYTHON_BINDINGS_H_
#define TENZR_SRC_PYTHON_BINDINGS_H_

#include <nanobind/nanobind.h>
#include <nanobind/ndarray.h>
#include <nanobind/operators.h>
#include <nanobind/stl/function.h>
#include <nanobind/stl/optional.h>
#include <nanobind/stl/pair.h>
#include <nanobind/stl/string.h>
#include <nanobind/stl/string_view.h>
#include <nanobind/stl/tuple.h>
#include <nanobind/stl/vector.h>

#include "tenzr/modules/profiler/profiler.h"
#include "tenzr/modules/random.h"
#include "tenzr/numpy/numpy.h"
#include "tenzr/plax/plax.h"
#include "tenzr/runtime/hardware_manager.h"
#include "tenzr/scipy/scipy.h"
#include "tenzr/sharding/mesh.h"
#include "tenzr/sharding/named_sharding.h"
#include "tenzr/sharding/partition_spec.h"
#include "tenzr/sharding/shardy.h"
#include "tenzr/tenzr.h"
#include "tenzr/txpr/control_flow.h"
#include "tenzr/txpr/custom_vjp.h"
#include "tenzr/txpr/grad.h"
#include "tenzr/txpr/jit.h"
#include "tenzr/txpr/jvp.h"
#include "tenzr/txpr/tracer.h"
#include "tenzr/txpr/vmap.h"

namespace nb = nanobind;
using namespace nb::literals;

namespace tenzr::python {

  void BindArray(nb::module_& m);
  void BindTransforms(nb::module_& m);
  void BindPlax(nb::module_& m);
  void BindRandom(nb::module_& m);
  void BindSharding(nb::module_& m);
  void BindScipy(nb::module_& m);
  void BindProfiler(nb::module_& m);

  // DLPack dtype mapping (complete 30-type coverage; throws on unknown).
  inline nb::dlpack::dtype ToDlpackDtype(DType dt) {
    using nb::dlpack::dtype_code;
    nb::dlpack::dtype d;
    d.lanes = 1;
    switch (dt) {
      case DType::kBool:
        d.code = static_cast<uint8_t>(dtype_code::Bool);
        d.bits = 8;
        break;
      case DType::kInt8:
        d.code = static_cast<uint8_t>(dtype_code::Int);
        d.bits = 8;
        break;
      case DType::kUInt8:
        d.code = static_cast<uint8_t>(dtype_code::UInt);
        d.bits = 8;
        break;
      case DType::kInt16:
        d.code = static_cast<uint8_t>(dtype_code::Int);
        d.bits = 16;
        break;
      case DType::kUInt16:
        d.code = static_cast<uint8_t>(dtype_code::UInt);
        d.bits = 16;
        break;
      case DType::kInt32:
        d.code = static_cast<uint8_t>(dtype_code::Int);
        d.bits = 32;
        break;
      case DType::kUInt32:
        d.code = static_cast<uint8_t>(dtype_code::UInt);
        d.bits = 32;
        break;
      case DType::kInt64:
        d.code = static_cast<uint8_t>(dtype_code::Int);
        d.bits = 64;
        break;
      case DType::kUInt64:
        d.code = static_cast<uint8_t>(dtype_code::UInt);
        d.bits = 64;
        break;
      case DType::kFloat16:
        d.code = static_cast<uint8_t>(dtype_code::Float);
        d.bits = 16;
        break;
      case DType::kBFloat16:
        d.code = static_cast<uint8_t>(dtype_code::Bfloat);
        d.bits = 16;
        break;
      case DType::kFloat32:
        d.code = static_cast<uint8_t>(dtype_code::Float);
        d.bits = 32;
        break;
      case DType::kFloat64:
        d.code = static_cast<uint8_t>(dtype_code::Float);
        d.bits = 64;
        break;
      case DType::kFloat8E4M3FN:
        d.code = static_cast<uint8_t>(dtype_code::Float8_E4M3FN);
        d.bits = 8;
        break;
      case DType::kFloat8E5M2:
        d.code = static_cast<uint8_t>(dtype_code::Float8_E5M2);
        d.bits = 8;
        break;
      case DType::kFloat8E3M4:
        d.code = static_cast<uint8_t>(dtype_code::Float8_E3M4);
        d.bits = 8;
        break;
      case DType::kFloat8E4M3:
        d.code = static_cast<uint8_t>(dtype_code::Float8_E4M3);
        d.bits = 8;
        break;
      case DType::kFloat8E4M3FNUZ:
        d.code = static_cast<uint8_t>(dtype_code::Float8_E4M3FNUZ);
        d.bits = 8;
        break;
      case DType::kFloat8E4M3B11FNUZ:
        d.code = static_cast<uint8_t>(dtype_code::Float8_E4M3B11FNUZ);
        d.bits = 8;
        break;
      case DType::kFloat8E5M2FNUZ:
        d.code = static_cast<uint8_t>(dtype_code::Float8_E5M2FNUZ);
        d.bits = 8;
        break;
      case DType::kFloat8E8M0FNU:
        d.code = static_cast<uint8_t>(dtype_code::Float8_E8M0FNU);
        d.bits = 8;
        break;
      case DType::kFloat6E2M3FN:
        d.code = static_cast<uint8_t>(dtype_code::Float6_E2M3FN);
        d.bits = 6;
        break;
      case DType::kFloat6E3M2FN:
        d.code = static_cast<uint8_t>(dtype_code::Float6_E3M2FN);
        d.bits = 6;
        break;
      case DType::kFloat4E2M1FN:
        d.code = static_cast<uint8_t>(dtype_code::Float4_E2M1FN);
        d.bits = 4;
        break;
      case DType::kInt1:
        d.code = static_cast<uint8_t>(dtype_code::Int);
        d.bits = 1;
        break;
      case DType::kUInt1:
        d.code = static_cast<uint8_t>(dtype_code::UInt);
        d.bits = 1;
        break;
      case DType::kInt2:
        d.code = static_cast<uint8_t>(dtype_code::Int);
        d.bits = 2;
        break;
      case DType::kUInt2:
        d.code = static_cast<uint8_t>(dtype_code::UInt);
        d.bits = 2;
        break;
      case DType::kInt4:
        d.code = static_cast<uint8_t>(dtype_code::Int);
        d.bits = 4;
        break;
      case DType::kUInt4:
        d.code = static_cast<uint8_t>(dtype_code::UInt);
        d.bits = 4;
        break;
      default:
        TENZR_THROW("Unsupported DType for DLPack export: " + std::string(tenzr::ToString(dt)));
    }
    return d;
  }

  inline DType FromDlpackDtype(nb::dlpack::dtype d) {
    using nb::dlpack::dtype_code;
    const auto code = static_cast<uint8_t>(d.code);
    if (code == static_cast<uint8_t>(dtype_code::Bool)) {
      return DType::kBool;
    }
    if (code == static_cast<uint8_t>(dtype_code::Float)) {
      if (d.bits == 16) {
        return DType::kFloat16;
      }
      if (d.bits == 32) {
        return DType::kFloat32;
      }
      if (d.bits == 64) {
        return DType::kFloat64;
      }
    }
    if (code == static_cast<uint8_t>(dtype_code::Bfloat) && d.bits == 16) {
      return DType::kBFloat16;
    }
    if (code == static_cast<uint8_t>(dtype_code::Int)) {
      if (d.bits == 8) {
        return DType::kInt8;
      }
      if (d.bits == 16) {
        return DType::kInt16;
      }
      if (d.bits == 32) {
        return DType::kInt32;
      }
      if (d.bits == 64) {
        return DType::kInt64;
      }
      if (d.bits == 4) {
        return DType::kInt4;
      }
      if (d.bits == 2) {
        return DType::kInt2;
      }
      if (d.bits == 1) {
        return DType::kInt1;
      }
    }
    if (code == static_cast<uint8_t>(dtype_code::UInt)) {
      if (d.bits == 8) {
        return DType::kUInt8;
      }
      if (d.bits == 16) {
        return DType::kUInt16;
      }
      if (d.bits == 32) {
        return DType::kUInt32;
      }
      if (d.bits == 64) {
        return DType::kUInt64;
      }
      if (d.bits == 4) {
        return DType::kUInt4;
      }
      if (d.bits == 2) {
        return DType::kUInt2;
      }
      if (d.bits == 1) {
        return DType::kUInt1;
      }
    }
    if (code == static_cast<uint8_t>(dtype_code::Float8_E4M3FN)) {
      return DType::kFloat8E4M3FN;
    }
    if (code == static_cast<uint8_t>(dtype_code::Float8_E5M2)) {
      return DType::kFloat8E5M2;
    }
    if (code == static_cast<uint8_t>(dtype_code::Float8_E3M4)) {
      return DType::kFloat8E3M4;
    }
    if (code == static_cast<uint8_t>(dtype_code::Float8_E4M3)) {
      return DType::kFloat8E4M3;
    }
    if (code == static_cast<uint8_t>(dtype_code::Float8_E4M3FNUZ)) {
      return DType::kFloat8E4M3FNUZ;
    }
    if (code == static_cast<uint8_t>(dtype_code::Float8_E4M3B11FNUZ)) {
      return DType::kFloat8E4M3B11FNUZ;
    }
    if (code == static_cast<uint8_t>(dtype_code::Float8_E5M2FNUZ)) {
      return DType::kFloat8E5M2FNUZ;
    }
    if (code == static_cast<uint8_t>(dtype_code::Float8_E8M0FNU)) {
      return DType::kFloat8E8M0FNU;
    }
    if (code == static_cast<uint8_t>(dtype_code::Float6_E2M3FN)) {
      return DType::kFloat6E2M3FN;
    }
    if (code == static_cast<uint8_t>(dtype_code::Float6_E3M2FN)) {
      return DType::kFloat6E3M2FN;
    }
    if (code == static_cast<uint8_t>(dtype_code::Float4_E2M1FN)) {
      return DType::kFloat4E2M1FN;
    }
    TENZR_THROW("Unsupported DLPack dtype (code=" + std::to_string(static_cast<int>(d.code)) + ", bits=" + std::to_string(static_cast<int>(d.bits)) + ")");
  }

  inline double ReadScalarAsDouble(const void* ptr, DType dt) {
    switch (dt) {
      case DType::kBool:
        return *reinterpret_cast<const uint8_t*>(ptr) != 0 ? 1.0 : 0.0;
      case DType::kInt8:
        return static_cast<double>(*reinterpret_cast<const int8_t*>(ptr));
      case DType::kUInt8:
        return static_cast<double>(*reinterpret_cast<const uint8_t*>(ptr));
      case DType::kInt16:
        return static_cast<double>(*reinterpret_cast<const int16_t*>(ptr));
      case DType::kUInt16:
        return static_cast<double>(*reinterpret_cast<const uint16_t*>(ptr));
      case DType::kInt32:
        return static_cast<double>(*reinterpret_cast<const int32_t*>(ptr));
      case DType::kUInt32:
        return static_cast<double>(*reinterpret_cast<const uint32_t*>(ptr));
      case DType::kInt64:
        return static_cast<double>(*reinterpret_cast<const int64_t*>(ptr));
      case DType::kUInt64:
        return static_cast<double>(*reinterpret_cast<const uint64_t*>(ptr));
      case DType::kFloat16:
        return static_cast<double>(*reinterpret_cast<const tenzr::float16*>(ptr));
      case DType::kBFloat16:
        return static_cast<double>(*reinterpret_cast<const tenzr::bfloat16*>(ptr));
      case DType::kFloat32:
        return static_cast<double>(*reinterpret_cast<const float*>(ptr));
      case DType::kFloat64:
        return *reinterpret_cast<const double*>(ptr);
      case DType::kFloat8E4M3FN:
        return static_cast<double>(*reinterpret_cast<const tenzr::float8_e4m3fn*>(ptr));
      case DType::kFloat8E5M2:
        return static_cast<double>(*reinterpret_cast<const tenzr::float8_e5m2*>(ptr));
      case DType::kFloat8E3M4:
        return static_cast<double>(*reinterpret_cast<const tenzr::float8_e3m4*>(ptr));
      case DType::kFloat8E4M3:
        return static_cast<double>(*reinterpret_cast<const tenzr::float8_e4m3*>(ptr));
      case DType::kFloat8E4M3FNUZ:
        return static_cast<double>(*reinterpret_cast<const tenzr::float8_e4m3fnuz*>(ptr));
      case DType::kFloat8E4M3B11FNUZ:
        return static_cast<double>(*reinterpret_cast<const tenzr::float8_e4m3b11fnuz*>(ptr));
      case DType::kFloat8E5M2FNUZ:
        return static_cast<double>(*reinterpret_cast<const tenzr::float8_e5m2fnuz*>(ptr));
      case DType::kFloat8E8M0FNU:
        return static_cast<double>(*reinterpret_cast<const tenzr::float8_e8m0fnu*>(ptr));
      case DType::kFloat6E2M3FN:
        return static_cast<double>(*reinterpret_cast<const tenzr::float6_e2m3fn*>(ptr));
      case DType::kFloat6E3M2FN:
        return static_cast<double>(*reinterpret_cast<const tenzr::float6_e3m2fn*>(ptr));
      case DType::kFloat4E2M1FN:
        return static_cast<double>(*reinterpret_cast<const tenzr::float4_e2m1fn*>(ptr));
      case DType::kInt1:
        return static_cast<double>(static_cast<int64_t>(*reinterpret_cast<const tenzr::int1*>(ptr)));
      case DType::kUInt1:
        return static_cast<double>(static_cast<int64_t>(*reinterpret_cast<const tenzr::uint1*>(ptr)));
      case DType::kInt2:
        return static_cast<double>(static_cast<int64_t>(*reinterpret_cast<const tenzr::int2*>(ptr)));
      case DType::kUInt2:
        return static_cast<double>(static_cast<int64_t>(*reinterpret_cast<const tenzr::uint2*>(ptr)));
      case DType::kInt4:
        return static_cast<double>(static_cast<int64_t>(*reinterpret_cast<const tenzr::int4*>(ptr)));
      case DType::kUInt4:
        return static_cast<double>(static_cast<int64_t>(*reinterpret_cast<const tenzr::uint4*>(ptr)));
      default:
        TENZR_THROW("Unsupported DType in read_scalar_as_double");
    }
  }

  inline DType PythonScalarDtype(nb::handle obj) {
    if (nb::isinstance<nb::bool_>(obj)) {
      return DType::kBool;
    }
    if (nb::isinstance<nb::int_>(obj)) {
      return DType::kInt32;
    }
    return DType::kFloat32;
  }

  inline DType DefaultScalarDtype(nb::handle obj) {
    if (nb::isinstance<nb::bool_>(obj)) {
      return DType::kBool;
    }
    if (nb::isinstance<nb::int_>(obj)) {
      return DType::kInt64;
    }
    return DType::kFloat64;
  }

}  // namespace tenzr::python

#endif  // TENZR_SRC_PYTHON_BINDINGS_H_