tensoracle 1.6.0__tar.gz → 2.1.0__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (211) hide show
  1. tensoracle-2.1.0/CMakeLists.txt +36 -0
  2. tensoracle-2.1.0/PKG-INFO +179 -0
  3. tensoracle-2.1.0/README.md +147 -0
  4. tensoracle-2.1.0/bindings.cpp +213 -0
  5. tensoracle-2.1.0/build_backend.py +132 -0
  6. tensoracle-2.1.0/cpp/CMakeLists.txt +123 -0
  7. tensoracle-2.1.0/cpp/cmake/TensoracleDependencies.cmake +21 -0
  8. tensoracle-2.1.0/cpp/cmake/TensoracleJson.cmake +12 -0
  9. tensoracle-2.1.0/cpp/cmake/tensoracleConfig.cmake.in +4 -0
  10. tensoracle-2.1.0/cpp/include/tensoracle/builder.h +122 -0
  11. tensoracle-2.1.0/cpp/include/tensoracle/builder.hpp +126 -0
  12. tensoracle-2.1.0/cpp/include/tensoracle/cache.h +110 -0
  13. tensoracle-2.1.0/cpp/include/tensoracle/cache.hpp +57 -0
  14. tensoracle-2.1.0/cpp/include/tensoracle/cubin.h +88 -0
  15. tensoracle-2.1.0/cpp/include/tensoracle/cubin.hpp +61 -0
  16. {tensoracle-1.6.0 → tensoracle-2.1.0}/cpp/include/tensoracle/cubin_plugin.hpp +0 -2
  17. tensoracle-2.1.0/cpp/include/tensoracle/engine.h +89 -0
  18. tensoracle-2.1.0/cpp/include/tensoracle/engine.hpp +94 -0
  19. tensoracle-2.1.0/cpp/include/tensoracle/plugins.h +30 -0
  20. tensoracle-2.1.0/cpp/include/tensoracle/plugins.hpp +29 -0
  21. tensoracle-2.1.0/cpp/include/tensoracle/tensoracle.h +163 -0
  22. tensoracle-2.1.0/cpp/include/tensoracle/tensoracle.hpp +75 -0
  23. tensoracle-2.1.0/cpp/src/abi.cpp +60 -0
  24. tensoracle-2.1.0/cpp/src/builder.cpp +213 -0
  25. tensoracle-2.1.0/cpp/src/builder_internal.hpp +14 -0
  26. tensoracle-2.1.0/cpp/src/bundle.cpp +330 -0
  27. tensoracle-2.1.0/cpp/src/bundle_internal.hpp +36 -0
  28. tensoracle-2.1.0/cpp/src/cache.cpp +248 -0
  29. tensoracle-2.1.0/cpp/src/cache.hpp +78 -0
  30. tensoracle-2.1.0/cpp/src/cache_management.cpp +117 -0
  31. tensoracle-2.1.0/cpp/src/cache_management_internal.cpp +198 -0
  32. tensoracle-2.1.0/cpp/src/cache_management_internal.hpp +45 -0
  33. tensoracle-2.1.0/cpp/src/cache_open.cpp +171 -0
  34. tensoracle-2.1.0/cpp/src/comparison.cpp +218 -0
  35. {tensoracle-1.6.0 → tensoracle-2.1.0}/cpp/src/cubin_plugin.cpp +4 -71
  36. tensoracle-2.1.0/cpp/src/engine.cpp +193 -0
  37. tensoracle-2.1.0/cpp/src/engine_internal.hpp +71 -0
  38. tensoracle-2.1.0/cpp/src/error.cpp +47 -0
  39. tensoracle-2.1.0/cpp/src/error.hpp +34 -0
  40. tensoracle-2.1.0/cpp/src/native_exports.hpp +112 -0
  41. tensoracle-2.1.0/cpp/src/plugins.cpp +99 -0
  42. tensoracle-2.1.0/cpp/src/producer.cpp +513 -0
  43. tensoracle-2.1.0/cpp/src/safetensors.cpp +124 -0
  44. tensoracle-2.1.0/cpp/src/safetensors.hpp +34 -0
  45. tensoracle-2.1.0/cpp/src/sha256.hpp +195 -0
  46. tensoracle-2.1.0/cpp/src/structure.hpp +28 -0
  47. tensoracle-2.1.0/cpp/src/tensoracle.map +50 -0
  48. tensoracle-2.1.0/cpp/tests/CMakeLists.txt +134 -0
  49. tensoracle-2.1.0/cpp/tests/abi.cpp +112 -0
  50. tensoracle-2.1.0/cpp/tests/builder.cpp +202 -0
  51. tensoracle-2.1.0/cpp/tests/builder_fixture.hpp +112 -0
  52. tensoracle-2.1.0/cpp/tests/builder_probe.cpp +26 -0
  53. tensoracle-2.1.0/cpp/tests/c_consumer.c +76 -0
  54. tensoracle-2.1.0/cpp/tests/cache.cpp +263 -0
  55. tensoracle-2.1.0/cpp/tests/cache_management.cpp +316 -0
  56. tensoracle-2.1.0/cpp/tests/cache_probe.cpp +21 -0
  57. tensoracle-2.1.0/cpp/tests/comparison.cpp +199 -0
  58. tensoracle-2.1.0/cpp/tests/consumer/CMakeLists.txt +12 -0
  59. tensoracle-2.1.0/cpp/tests/consumer/main.cpp +43 -0
  60. tensoracle-2.1.0/cpp/tests/cubin.cpp +112 -0
  61. tensoracle-2.1.0/cpp/tests/driver_failures.cpp +84 -0
  62. tensoracle-2.1.0/cpp/tests/execution.cpp +420 -0
  63. tensoracle-2.1.0/cpp/tests/installed.cmake +15 -0
  64. tensoracle-2.1.0/cpp/tests/native_exports.cpp +33 -0
  65. tensoracle-2.1.0/cpp/tests/native_exports_fixture.cpp +12 -0
  66. tensoracle-2.1.0/cpp/tests/plugin_fixture.cpp +30 -0
  67. tensoracle-2.1.0/cpp/tests/plugins.cpp +93 -0
  68. tensoracle-2.1.0/cpp/tests/precision.cpp +36 -0
  69. tensoracle-2.1.0/cpp/tests/utilities.cpp +64 -0
  70. tensoracle-2.1.0/cpp/tests/weights.cpp +344 -0
  71. tensoracle-2.1.0/cubin_bindings.cpp +161 -0
  72. tensoracle-2.1.0/docs/api/bundle.md +23 -0
  73. tensoracle-2.1.0/docs/api/index.md +32 -0
  74. tensoracle-2.1.0/docs/api/producer.md +35 -0
  75. tensoracle-2.1.0/docs/cpp-python-repositories.md +86 -0
  76. tensoracle-2.1.0/docs/index.md +273 -0
  77. tensoracle-2.1.0/docs/migration.md +62 -0
  78. {tensoracle-1.6.0 → tensoracle-2.1.0}/docs/optimization.md +4 -4
  79. tensoracle-2.1.0/docs/precision.md +141 -0
  80. tensoracle-2.1.0/docs/tensorrt-sdk.md +125 -0
  81. tensoracle-2.1.0/examples/cubin_plugin/CMakeLists.txt +48 -0
  82. tensoracle-2.1.0/examples/cubin_plugin/README.md +95 -0
  83. tensoracle-2.1.0/examples/cubin_plugin/bindings.cpp +165 -0
  84. tensoracle-2.1.0/examples/cubin_plugin/bundle/bundle.cpp +5199 -0
  85. {tensoracle-1.6.0 → tensoracle-2.1.0}/examples/cubin_plugin/bundle/bundle.json +172 -274
  86. tensoracle-2.1.0/examples/cubin_plugin/bundle/cubins/sm120/bias.cubin +0 -0
  87. tensoracle-2.1.0/examples/cubin_plugin/bundle/cubins/sm120/dot.specialized.cubin +0 -0
  88. tensoracle-2.1.0/examples/cubin_plugin/bundle/cubins/sm120/dot.tf32-control.cubin +0 -0
  89. tensoracle-2.1.0/examples/cubin_plugin/bundle/cubins/sm120/dot.unspecialized.cubin +0 -0
  90. tensoracle-2.1.0/examples/cubin_plugin/bundle/cubins/sm120/scale.cubin +0 -0
  91. tensoracle-2.1.0/examples/cubin_plugin/bundle/cubins/sm86/bias.cubin +0 -0
  92. tensoracle-2.1.0/examples/cubin_plugin/bundle/cubins/sm86/dot.specialized.cubin +0 -0
  93. tensoracle-2.1.0/examples/cubin_plugin/bundle/cubins/sm86/dot.tf32-control.cubin +0 -0
  94. tensoracle-2.1.0/examples/cubin_plugin/bundle/cubins/sm86/dot.unspecialized.cubin +0 -0
  95. tensoracle-2.1.0/examples/cubin_plugin/bundle/cubins/sm86/scale.cubin +0 -0
  96. tensoracle-2.1.0/examples/cubin_plugin/bundle/cubins/sm89/bias.cubin +0 -0
  97. tensoracle-2.1.0/examples/cubin_plugin/bundle/cubins/sm89/dot.specialized.cubin +0 -0
  98. tensoracle-2.1.0/examples/cubin_plugin/bundle/cubins/sm89/dot.tf32-control.cubin +0 -0
  99. tensoracle-2.1.0/examples/cubin_plugin/bundle/cubins/sm89/dot.unspecialized.cubin +0 -0
  100. tensoracle-2.1.0/examples/cubin_plugin/bundle/cubins/sm89/scale.cubin +0 -0
  101. tensoracle-2.1.0/examples/cubin_plugin/example_cubin_plugin/__init__.py +36 -0
  102. tensoracle-2.1.0/examples/cubin_plugin/kernels.cu +40 -0
  103. tensoracle-2.1.0/examples/cubin_plugin/kernels.py +44 -0
  104. {tensoracle-1.6.0 → tensoracle-2.1.0}/examples/cubin_plugin/pyproject.toml +9 -1
  105. tensoracle-2.1.0/examples/wheel_consumer/CMakeLists.txt +38 -0
  106. tensoracle-2.1.0/examples/wheel_consumer/README.md +69 -0
  107. tensoracle-2.1.0/examples/wheel_consumer/bindings.cpp +124 -0
  108. tensoracle-2.1.0/examples/wheel_consumer/pyproject.toml +22 -0
  109. tensoracle-2.1.0/examples/wheel_consumer/src/tensoracle_consumer/__init__.py +5 -0
  110. {tensoracle-1.6.0 → tensoracle-2.1.0}/mkdocs.yml +3 -2
  111. {tensoracle-1.6.0 → tensoracle-2.1.0}/pyproject.toml +32 -21
  112. tensoracle-2.1.0/src/tensoracle/__init__.py +16 -0
  113. tensoracle-2.1.0/src/tensoracle/_cli.py +172 -0
  114. tensoracle-2.1.0/src/tensoracle/_comparison.py +46 -0
  115. {tensoracle-1.6.0 → tensoracle-2.1.0}/src/tensoracle/_version.py +2 -2
  116. tensoracle-2.1.0/src/tensoracle/bundle.py +45 -0
  117. tensoracle-2.1.0/src/tensoracle/producer.py +48 -0
  118. {tensoracle-1.6.0 → tensoracle-2.1.0}/tests/native/CMakeLists.txt +0 -19
  119. tensoracle-2.1.0/tests/native/cache/CMakeLists.txt +11 -0
  120. tensoracle-2.1.0/tests/native/cache/writer.cpp +18 -0
  121. {tensoracle-1.6.0 → tensoracle-2.1.0}/tests/native_build.py +0 -14
  122. {tensoracle-1.6.0 → tensoracle-2.1.0}/tests/test_bundle.py +8 -12
  123. {tensoracle-1.6.0 → tensoracle-2.1.0}/tests/test_cache_cli.py +186 -60
  124. tensoracle-2.1.0/tests/test_comparison_binding.py +103 -0
  125. {tensoracle-1.6.0 → tensoracle-2.1.0}/tests/test_convert.py +3 -3
  126. {tensoracle-1.6.0 → tensoracle-2.1.0}/tests/test_cpp_core.py +25 -2
  127. {tensoracle-1.6.0 → tensoracle-2.1.0}/tests/test_distribution.py +304 -21
  128. {tensoracle-1.6.0 → tensoracle-2.1.0}/tests/test_example_plugin.py +107 -89
  129. tensoracle-2.1.0/tests/test_native.py +26 -0
  130. tensoracle-2.1.0/tests/test_plugin_io.py +34 -0
  131. tensoracle-2.1.0/tests/test_producer.py +230 -0
  132. tensoracle-2.1.0/tests/test_public_api.py +56 -0
  133. {tensoracle-1.6.0 → tensoracle-2.1.0}/uv.lock +27 -28
  134. tensoracle-1.6.0/PKG-INFO +0 -491
  135. tensoracle-1.6.0/README.md +0 -458
  136. tensoracle-1.6.0/build_backend.py +0 -94
  137. tensoracle-1.6.0/cpp/CMakeLists.txt +0 -28
  138. tensoracle-1.6.0/docs/api/bundle.md +0 -21
  139. tensoracle-1.6.0/docs/api/index.md +0 -49
  140. tensoracle-1.6.0/docs/api/producer.md +0 -11
  141. tensoracle-1.6.0/docs/cpp-python-repositories.md +0 -65
  142. tensoracle-1.6.0/docs/index.md +0 -296
  143. tensoracle-1.6.0/docs/precision.md +0 -154
  144. tensoracle-1.6.0/docs/tensorrt-sdk.md +0 -215
  145. tensoracle-1.6.0/examples/cubin_plugin/CMakeLists.txt +0 -28
  146. tensoracle-1.6.0/examples/cubin_plugin/README.md +0 -50
  147. tensoracle-1.6.0/examples/cubin_plugin/bundle/bundle.cpp +0 -5647
  148. tensoracle-1.6.0/examples/cubin_plugin/bundle/cubins/sm120/bias.cubin +0 -0
  149. tensoracle-1.6.0/examples/cubin_plugin/bundle/cubins/sm120/dot.specialized.cubin +0 -0
  150. tensoracle-1.6.0/examples/cubin_plugin/bundle/cubins/sm120/dot.tf32-control.cubin +0 -0
  151. tensoracle-1.6.0/examples/cubin_plugin/bundle/cubins/sm120/dot.unspecialized.cubin +0 -0
  152. tensoracle-1.6.0/examples/cubin_plugin/bundle/cubins/sm120/scale.cubin +0 -0
  153. tensoracle-1.6.0/examples/cubin_plugin/bundle/cubins/sm86/bias.cubin +0 -0
  154. tensoracle-1.6.0/examples/cubin_plugin/bundle/cubins/sm86/dot.specialized.cubin +0 -0
  155. tensoracle-1.6.0/examples/cubin_plugin/bundle/cubins/sm86/dot.tf32-control.cubin +0 -0
  156. tensoracle-1.6.0/examples/cubin_plugin/bundle/cubins/sm86/dot.unspecialized.cubin +0 -0
  157. tensoracle-1.6.0/examples/cubin_plugin/bundle/cubins/sm86/scale.cubin +0 -0
  158. tensoracle-1.6.0/examples/cubin_plugin/bundle/cubins/sm89/bias.cubin +0 -0
  159. tensoracle-1.6.0/examples/cubin_plugin/bundle/cubins/sm89/dot.specialized.cubin +0 -0
  160. tensoracle-1.6.0/examples/cubin_plugin/bundle/cubins/sm89/dot.tf32-control.cubin +0 -0
  161. tensoracle-1.6.0/examples/cubin_plugin/bundle/cubins/sm89/dot.unspecialized.cubin +0 -0
  162. tensoracle-1.6.0/examples/cubin_plugin/bundle/cubins/sm89/scale.cubin +0 -0
  163. tensoracle-1.6.0/examples/cubin_plugin/example_cubin_plugin/__init__.py +0 -89
  164. tensoracle-1.6.0/examples/cubin_plugin/kernels.py +0 -92
  165. tensoracle-1.6.0/src/tensoracle/__init__.py +0 -60
  166. tensoracle-1.6.0/src/tensoracle/_build.py +0 -172
  167. tensoracle-1.6.0/src/tensoracle/_cache.py +0 -375
  168. tensoracle-1.6.0/src/tensoracle/_cache_format.py +0 -78
  169. tensoracle-1.6.0/src/tensoracle/_cli.py +0 -250
  170. tensoracle-1.6.0/src/tensoracle/_comparison.py +0 -117
  171. tensoracle-1.6.0/src/tensoracle/_execution.py +0 -187
  172. tensoracle-1.6.0/src/tensoracle/_memory.py +0 -18
  173. tensoracle-1.6.0/src/tensoracle/_plugins.py +0 -223
  174. tensoracle-1.6.0/src/tensoracle/bundle.py +0 -221
  175. tensoracle-1.6.0/src/tensoracle/producer.py +0 -265
  176. tensoracle-1.6.0/tests/build_memory_runtime.py +0 -89
  177. tensoracle-1.6.0/tests/cubin_io/kernels.cu +0 -17
  178. tensoracle-1.6.0/tests/cubin_io/plugin.cpp +0 -106
  179. tensoracle-1.6.0/tests/example_runtime.py +0 -299
  180. tensoracle-1.6.0/tests/plugin_io_runtime.py +0 -35
  181. tensoracle-1.6.0/tests/precision_runtime.py +0 -164
  182. tensoracle-1.6.0/tests/test_build.py +0 -263
  183. tensoracle-1.6.0/tests/test_build_memory.py +0 -93
  184. tensoracle-1.6.0/tests/test_cache.py +0 -809
  185. tensoracle-1.6.0/tests/test_cache_management.py +0 -198
  186. tensoracle-1.6.0/tests/test_cuda_graphs.py +0 -191
  187. tensoracle-1.6.0/tests/test_engine_io.py +0 -203
  188. tensoracle-1.6.0/tests/test_execution.py +0 -572
  189. tensoracle-1.6.0/tests/test_gpu_cache.py +0 -204
  190. tensoracle-1.6.0/tests/test_plugin_io.py +0 -192
  191. tensoracle-1.6.0/tests/test_plugins.py +0 -163
  192. tensoracle-1.6.0/tests/test_precision.py +0 -107
  193. tensoracle-1.6.0/tests/test_producer.py +0 -279
  194. tensoracle-1.6.0/tests/test_shape_tensor_outputs.py +0 -281
  195. tensoracle-1.6.0/tests/test_weight_hash.py +0 -339
  196. {tensoracle-1.6.0 → tensoracle-2.1.0}/.gitignore +0 -0
  197. {tensoracle-1.6.0 → tensoracle-2.1.0}/LICENSE +0 -0
  198. {tensoracle-1.6.0 → tensoracle-2.1.0}/cpp/cmake/TensoracleCubinPlugin.cmake +0 -0
  199. {tensoracle-1.6.0 → tensoracle-2.1.0}/cpp/cmake/TensoracleCuda.cmake +0 -0
  200. {tensoracle-1.6.0 → tensoracle-2.1.0}/cpp/tests/core.cpp +0 -0
  201. {tensoracle-1.6.0 → tensoracle-2.1.0}/examples/cubin_plugin/example_cubin_plugin/weights.safetensors +0 -0
  202. {tensoracle-1.6.0 → tensoracle-2.1.0}/examples/cubin_plugin/plugin.cpp +0 -0
  203. {tensoracle-1.6.0 → tensoracle-2.1.0}/src/tensoracle/convert.py +0 -0
  204. {tensoracle-1.6.0 → tensoracle-2.1.0}/tests/cubin_io/check.cpp +0 -0
  205. {tensoracle-1.6.0 → tensoracle-2.1.0}/tests/cubins/ptxas-12.9.86/sm120.cubin +0 -0
  206. {tensoracle-1.6.0 → tensoracle-2.1.0}/tests/cubins/ptxas-12.9.86/sm86.cubin +0 -0
  207. {tensoracle-1.6.0 → tensoracle-2.1.0}/tests/cubins/ptxas-12.9.86/sm89.cubin +0 -0
  208. {tensoracle-1.6.0 → tensoracle-2.1.0}/tests/native/export.h +0 -0
  209. {tensoracle-1.6.0 → tensoracle-2.1.0}/tests/test_cmake_cuda.py +0 -0
  210. {tensoracle-1.6.0 → tensoracle-2.1.0}/tests/test_cmake_plugin.py +0 -0
  211. {tensoracle-1.6.0 → tensoracle-2.1.0}/tests/test_comparison.py +0 -0
@@ -0,0 +1,36 @@
1
+ cmake_minimum_required(VERSION 3.26)
2
+ project(tensoracle_python LANGUAGES CXX)
3
+
4
+ find_package(Python 3.12 REQUIRED COMPONENTS Interpreter Development.Module Development.SABIModule)
5
+ find_package(nanobind CONFIG REQUIRED)
6
+
7
+ # The build environment supplies headers and the link library from wheels.
8
+ # The standalone C++ project receives paths and never discovers Python itself.
9
+ execute_process(COMMAND "${Python_EXECUTABLE}" -c
10
+ "from importlib.metadata import distribution; print(distribution('tensorrt-cu13-libs').locate_file('tensorrt_libs')); print(distribution('nvidia-cuda-runtime').locate_file('nvidia/cu13'))"
11
+ OUTPUT_VARIABLE _dependencies OUTPUT_STRIP_TRAILING_WHITESPACE COMMAND_ERROR_IS_FATAL ANY)
12
+ string(REPLACE "\n" ";" _dependencies "${_dependencies}")
13
+ list(GET _dependencies 0 _tensorrt_libs)
14
+ list(GET _dependencies 1 _cuda_root)
15
+ find_library(_nvinfer NAMES nvinfer nvinfer_11 libnvinfer.so.11
16
+ PATHS "${_tensorrt_libs}" NO_DEFAULT_PATH REQUIRED)
17
+ set(TENSORACLE_TENSORRT_LIBRARY "${_nvinfer}" CACHE FILEPATH "TensorRT from the build wheel" FORCE)
18
+ set(CUDAToolkit_ROOT "${_cuda_root}" CACHE PATH "CUDA headers from the build wheel" FORCE)
19
+ set(TENSORACLE_INSTALL_ROOT "tensoracle/cpp" CACHE STRING "Wheel SDK prefix" FORCE)
20
+ set(TENSORACLE_INSTALL_RPATH "$ORIGIN/../../../tensorrt_libs" CACHE STRING "TensorRT wheel lookup" FORCE)
21
+ set(CMAKE_INSTALL_LIBDIR lib CACHE STRING "Wheel library directory" FORCE)
22
+ set(BUILD_TESTING OFF CACHE BOOL "Tests are built separately" FORCE)
23
+ set(TENSORACLE_BUILD_CUBIN_HELPER OFF CACHE BOOL "Ship the source-based helper" FORCE)
24
+
25
+ # A staged sdist has cpp next to CMakeLists.txt; the normal checkout has ../cpp.
26
+ if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/cpp/CMakeLists.txt")
27
+ set(_cpp "${CMAKE_CURRENT_SOURCE_DIR}/cpp")
28
+ else()
29
+ set(_cpp "${CMAKE_CURRENT_SOURCE_DIR}/../cpp")
30
+ endif()
31
+ add_subdirectory("${_cpp}" cpp)
32
+ nanobind_add_module(_native STABLE_ABI bindings.cpp cubin_bindings.cpp)
33
+ target_link_libraries(_native PRIVATE tensoracle::tensoracle)
34
+ set_target_properties(_native PROPERTIES
35
+ BUILD_RPATH_USE_ORIGIN YES INSTALL_RPATH "$ORIGIN/cpp/lib")
36
+ install(TARGETS _native LIBRARY DESTINATION tensoracle RUNTIME DESTINATION tensoracle)
@@ -0,0 +1,179 @@
1
+ Metadata-Version: 2.4
2
+ Name: tensoracle
3
+ Version: 2.1.0
4
+ Summary: Build, save, reuse, and execute TensorRT engines from safetensors weights
5
+ Author-Email: minamik <mia@sync.dev>
6
+ License-Expression: MIT
7
+ License-File: LICENSE
8
+ Classifier: Development Status :: 5 - Production/Stable
9
+ Classifier: Environment :: GPU :: NVIDIA CUDA :: 13
10
+ Classifier: Intended Audience :: Developers
11
+ Classifier: Operating System :: POSIX :: Linux
12
+ Classifier: Operating System :: Microsoft :: Windows
13
+ Classifier: Programming Language :: Python :: 3.12
14
+ Classifier: Programming Language :: Python :: 3.13
15
+ Classifier: Topic :: Software Development :: Libraries
16
+ Project-URL: documentation, https://sync-dev-org.github.io/tensoracle/
17
+ Project-URL: repository, https://github.com/sync-dev-org/tensoracle
18
+ Project-URL: issues, https://github.com/sync-dev-org/tensoracle/issues
19
+ Requires-Python: <3.14,>=3.12
20
+ Requires-Dist: numpy>=2.0
21
+ Requires-Dist: safetensors>=0.4
22
+ Requires-Dist: tensorrt-cu13<12,>=11.0
23
+ Provides-Extra: convert
24
+ Requires-Dist: torch>=2.6; extra == "convert"
25
+ Provides-Extra: producer
26
+ Requires-Dist: nvidia-cuda-nvrtc==13.2.*; extra == "producer"
27
+ Requires-Dist: nvidia-cuda-nvcc==13.2.*; extra == "producer"
28
+ Requires-Dist: nvidia-cuda-cuobjdump==13.2.*; extra == "producer"
29
+ Requires-Dist: nvidia-cuda-crt==13.2.*; extra == "producer"
30
+ Requires-Dist: nvidia-nvvm==13.2.*; extra == "producer"
31
+ Description-Content-Type: text/markdown
32
+
33
+ # tensoracle
34
+
35
+ Build, save, reuse, and execute TensorRT engines through a C++ library embedded in your model package. Define the TensorRT network and engine-opening code in C++, then expose inference through your own binding, such as nanobind. Python utilities provide output comparison, cubin production and inspection, cache management commands, and `.pth` checkpoint conversion.
36
+
37
+ Engine building, reuse, execution, output comparison, and cubin production do not require torch. Only the optional converter requires `tensoracle[convert]`.
38
+
39
+ Read the [v2 user guide](https://sync-dev-org.github.io/tensoracle/) and [Migrate from 1.x](https://sync-dev-org.github.io/tensoracle/migration/).
40
+
41
+ ## Terms and workflows
42
+
43
+ - A **TensorRT engine** is TensorRT's executable form of a model, optimized for a GPU, precision, and TensorRT version. It can be held in memory or saved to disk.
44
+ - **Building an engine** means creating that executable form from the model package's C++ network-building function and safetensors weights.
45
+ - A **cubin plugin** is a native TensorRT plugin whose GPU kernels are supplied as compiled CUDA binary (`.cubin`) files. The model package owns its CUDA C++ kernels and plugin behavior.
46
+
47
+ | Role | Use | Additional requirements |
48
+ | --- | --- | --- |
49
+ | Inference application | Call the model package's inference binding. | NVIDIA GPU and driver; the model package manages CUDA buffers and streams. |
50
+ | Model package developer | Open and execute engines through tensoracle's C++ SDK, exposing inference through your own binding. | TensorRT 11 headers, CMake 3.26 or later, C++17 compiler, and a binding framework when needed. |
51
+ | Cubin plugin developer | Compile CUDA C++ kernels and package a native plugin with the model. | `tensoracle[producer]`, PyPI CUDA headers and tools, TensorRT SDK, CMake, and a C++17 compiler. |
52
+
53
+ v2.0.0 supplies a Linux x86_64 `manylinux_2_28`, `cp312-abi3` wheel, including WSL2. The package currently declares Python 3.12 and 3.13 support. Target GPUs are sm86, sm89, and sm120 with TensorRT 11; fp8 requires sm89 or sm120. Windows x86_64 users must keep using tensoracle 1.x until a later v2 minor release adds a Windows wheel.
54
+
55
+ The wheel contains the C++ public headers, `libtensoracle`, CMake package configuration, and cubin plugin helper sources. TensorRT is resolved through a wheel dependency; the NVIDIA driver is not bundled. Public C++ interfaces and saved engine files remain compatible within a major version.
56
+
57
+ ## Build and reopen an engine in C++
58
+
59
+ Include `tensoracle/builder.hpp` and link `tensoracle::tensoracle` from `find_package(tensoracle 2 CONFIG REQUIRED)`. Locate the wheel's SDK with:
60
+
61
+ ```bash
62
+ python -c 'from importlib.metadata import distribution; print(distribution("tensoracle").locate_file("tensoracle/cpp"))'
63
+ ```
64
+
65
+ Pass that directory in `CMAKE_PREFIX_PATH`. The model package supplies a callback taking `nvinfer1::INetworkDefinition&` and `const tensoracle::weights_view&`. tensoracle reads declared safetensors and keeps their values alive through the build.
66
+
67
+ This excerpt follows the [separate consumer wheel](https://github.com/sync-dev-org/tensoracle/tree/main/examples/wheel_consumer). It assumes the example's `build_network` function, four-element float32 `bias.safetensors`, and caller-owned current CUDA context. The [user guide](https://sync-dev-org.github.io/tensoracle/#build-and-reopen-an-engine) shows the callback; the complete example supplies allocation, error handling, and nanobind integration.
68
+
69
+ ```cpp
70
+ #include <tensoracle/builder.hpp>
71
+
72
+ tensoracle::builder_options options;
73
+ options.name = "wheel-consumer-add";
74
+ options.graph_version = "1";
75
+ options.weights = {"bias.safetensors"};
76
+ options.inputs = {{"input", {4}, {4}, {4}}};
77
+ options.io = {{"input", TENSORACLE_FLOAT32, {4}, true},
78
+ {"output", TENSORACLE_FLOAT32, {4}, false}};
79
+ options.workspace_limit = 64 * 1024 * 1024;
80
+ options.optimization_level = 0;
81
+ auto engine = tensoracle::open_engine(options, build_network);
82
+ ```
83
+
84
+ On a valid cache hit, `tensoracle::open_engine` skips the callback, including after a process restart. Engine identity includes weights, plugin bytes, name, graph version, input profiles, builder optimization level, GPU compatibility properties, OS, CPU architecture, and TensorRT version. Update the graph version when network code changes; use separate names or versions for different precisions. Tensor types in the network determine precision, and TF32 is disabled.
85
+
86
+ Leave `workspace_limit` unset to use TensorRT's default or supply a positive byte count. It does not enter the cache identity; set `options.force_rebuild = true` to rebuild with a changed limit. Saved files live under `XDG_CACHE_HOME/tensoracle`, or `~/.cache/tensoracle` when unset. The [user guide](https://sync-dev-org.github.io/tensoracle/#save-and-reuse-an-engine) describes integrity checks, errors, and weight digest reuse.
87
+
88
+ ## Execute through the model package
89
+
90
+ The C++ model package calls `engine.create_execution()`, sets ranged input shapes together, queries output shapes and types, and enqueues every I/O device address on its own CUDA stream. Each execution unit retains the engine and owns an independent TensorRT context. The caller owns allocation, transfers, stream ordering, and completion checks. tensoracle does not synchronize enqueue or release; wait before changing shapes, freeing buffers, or closing the unit.
91
+
92
+ The [execution guide](https://sync-dev-org.github.io/tensoracle/#execute-an-engine) follows the consumer wheel's C++ API usage. The model package exposes its inference function to Python. Python applications do not construct or execute engines through tensoracle's Python utilities.
93
+
94
+ For CUDA Graphs, warm up one enqueue after each shape change before capturing on a nondefault stream. The graph fixes shapes, addresses, and execution unit; keep resources alive until work completes and the graph is destroyed. The caller creates and recreates graphs. Engines with loops, conditionals, or host-synchronizing layers or plugins cannot be captured. See the [capture guide](https://sync-dev-org.github.io/tensoracle/#capture-and-replay-with-cuda-graphs).
95
+
96
+ ## Manage saved engines
97
+
98
+ ```bash
99
+ tensoracle cache list
100
+ tensoracle cache prune --unused-for 30 --dry-run
101
+ tensoracle cache prune --unused-for 30
102
+ tensoracle cache prune --max-size 10GiB --dry-run
103
+ ```
104
+
105
+ These commands use the same cache directory as the C++ engine opener. They inspect headers and metadata without a GPU, original weights, or plan deserialization. Listing reports conditions, size, recorded time, workspace limit, and known weight paths. Empty files are `building or leftover`; invalid headers, including 1.x engine files, are `unrecognized`. Pruning never deletes either category.
106
+
107
+ An engine file's modification time records its initial save or last successful open. It does not indicate whether inference continues. Failed loads and management commands do not change it. Timestamp write failures do not prevent a successful load; records that could not be written, or predate timestamp recording, may be stale.
108
+
109
+ Pruning requires exactly one of `--unused-for DAYS` and `--max-size SIZE`. It prints a plan and rebuild notice before deletion and separates removed, excluded, missing, and failed files in its result. `--dry-run` changes nothing. Run pruning during maintenance without concurrent builds or opens. The [management guide](https://sync-dev-org.github.io/tensoracle/#manage-saved-engines) describes sizes, selection, and exit statuses.
110
+
111
+ ## Build a cubin plugin
112
+
113
+ Install `tensoracle[producer]` to generate sm86, sm89, and sm120 cubins from CUDA C++ source. The C++ producer uses NVRTC, `ptxas >=13.2,<14`, and `cuobjdump` from PyPI CUDA 13.2 wheels, without a GPU or OS-installed CUDA Toolkit. Triton kernels are not accepted. See the [producer API](https://sync-dev-org.github.io/tensoracle/api/producer/) for `KernelSpec` and `compile_bundle`.
114
+
115
+ Build the native plugin with [the wheel's C++ sources and CMake helpers](https://sync-dev-org.github.io/tensoracle/cpp-python-repositories/) and a separately prepared [TensorRT SDK](https://sync-dev-org.github.io/tensoracle/tensorrt-sdk/). Export TensorRT's `getCreators` and list all required libraries in `tensoracle::builder_options::plugins`. tensoracle inspects and registers them before construction or deserialization and hashes their bytes into the engine identity. Explicit registration is also available through `tensoracle/plugins.hpp`.
116
+
117
+ Package the plugin library and cubin bundle in your model wheel. `tensoracle.bundle.inspect_bundle` and `inspect_wheel` verify the manifest, cubin digests and actual architectures, declared argument ABI, and the plugin's embedded cubins without loading it. The manifest records NVRTC and ptxas versions and kernel source digests. The [plugin guide](https://sync-dev-org.github.io/tensoracle/#build-a-cubin-plugin) describes tensor types, formats, and the common lifecycle helper.
118
+
119
+ Follow the [four-stage optimization workflow](https://sync-dev-org.github.io/tensoracle/optimization/) and [precision guide](https://sync-dev-org.github.io/tensoracle/precision/) when validating a model or plugin.
120
+
121
+ ## Convert a PyTorch checkpoint
122
+
123
+ Install the conversion extra in the environment where you prepare weights:
124
+
125
+ ```bash
126
+ pip install 'tensoracle[convert]'
127
+ python -m tensoracle.convert model.pth model.safetensors
128
+ ```
129
+
130
+ Build scripts can call the same converter:
131
+
132
+ ```python
133
+ from tensoracle.convert import pth_to_safetensors
134
+
135
+ pth_to_safetensors("model.pth", "model.safetensors")
136
+ ```
137
+
138
+ The input must be a mapping with string keys whose values are tensors or nested mappings. Nested keys are joined with `.`: `{"state_dict": {"layer.weight": tensor}}` becomes `state_dict.layer.weight`. The converter preserves each tensor's dtype, shape, and values, including shared or noncontiguous dense tensors. It rejects non-tensor metadata and duplicate flattened names before writing the destination. Weight renaming and model-specific preprocessing belong to the model package.
139
+
140
+ The converter loads weights on the CPU with `torch.load(..., weights_only=True)` and rejects checkpoint objects instead of falling back to unrestricted pickle loading. The extra requires torch 2.6 or later, which includes the fix for [CVE-2025-32434](https://github.com/pytorch/pytorch/security/advisories/GHSA-53q9-r3pm-6pq6). Torch is imported only when conversion is called; it is absent from the package's required dependencies and is not imported by engine building or reuse. Calling the converter without torch raises `ImportError` with the installation command for `tensoracle[convert]`.
141
+
142
+ ## Compare output arrays
143
+
144
+ Run the TensorRT engines in your application, then pass one CPU NumPy output array and its corresponding reference array to `compare_outputs`. The first argument is the candidate; the second is the reference. Both arrays must have the same nonempty shape and be floating point, integer, or boolean arrays of the same kind. The function does not run an engine or decide whether its output passes a quality limit.
145
+
146
+ For floating point arrays, `OutputComparison` contains the non-finite count for each side, the maximum absolute error, the relative L2 error (difference norm divided by reference norm), and cosine similarity. Supply both `atol` and `rtol` to also get the count and fraction of elements outside `abs(candidate - reference) <= atol + rtol * abs(reference)`. Choose these limits for each output in your application. Floating point calculations use float64, including comparisons between fp16 and fp32 arrays. A non-finite value on either side makes all floating point measurements `None`. A zero reference norm makes relative L2 error `None`; a zero norm on either side makes cosine similarity `None`.
147
+
148
+ First, compare a saved correct output with the output from an fp32 TensorRT engine. The application creates and fixes the correct output outside tensoracle and runs the engine itself:
149
+
150
+ ```python
151
+ import numpy as np
152
+
153
+ from tensoracle import compare_outputs
154
+
155
+ correct_output = np.array([1.0, 2.0], dtype=np.float32) # Fixed reference output
156
+ fp32_output = np.array([1.0001, 2.0002], dtype=np.float32) # Engine output on CPU
157
+ fp32_result = compare_outputs(fp32_output, correct_output, atol=0.001, rtol=0.0)
158
+ fp32_within_tolerance = (
159
+ fp32_result.candidate_nonfinite_count == 0
160
+ and fp32_result.reference_nonfinite_count == 0
161
+ and fp32_result.outside_tolerance_count == 0
162
+ )
163
+ ```
164
+
165
+ Then use the fp32 TensorRT engine output as the reference for an fp16 TensorRT engine output. The application chooses limits for whole-array closeness or direction as well as any elementwise tolerance:
166
+
167
+ ```python
168
+ fp16_output = np.array([1.0, 2.0], dtype=np.float16) # Engine output on CPU
169
+ fp16_result = compare_outputs(fp16_output, fp32_output, atol=0.001, rtol=0.001)
170
+ fp16_accepted = (
171
+ fp16_result.candidate_nonfinite_count == 0
172
+ and fp16_result.reference_nonfinite_count == 0
173
+ and fp16_result.outside_tolerance_count == 0
174
+ and fp16_result.relative_l2_error is not None
175
+ and fp16_result.relative_l2_error <= 0.01 # Application's limit
176
+ )
177
+ ```
178
+
179
+ When tolerances are omitted, `outside_tolerance_count` and `outside_tolerance_fraction` are `None`. For integer and boolean arrays, use `mismatch_count` and `mismatch_fraction`; these measure exact equality instead of floating point distances. A zero mismatch count means those arrays agree. Shape, kind, or empty-array errors raise `ValueError`.
@@ -0,0 +1,147 @@
1
+ # tensoracle
2
+
3
+ Build, save, reuse, and execute TensorRT engines through a C++ library embedded in your model package. Define the TensorRT network and engine-opening code in C++, then expose inference through your own binding, such as nanobind. Python utilities provide output comparison, cubin production and inspection, cache management commands, and `.pth` checkpoint conversion.
4
+
5
+ Engine building, reuse, execution, output comparison, and cubin production do not require torch. Only the optional converter requires `tensoracle[convert]`.
6
+
7
+ Read the [v2 user guide](https://sync-dev-org.github.io/tensoracle/) and [Migrate from 1.x](https://sync-dev-org.github.io/tensoracle/migration/).
8
+
9
+ ## Terms and workflows
10
+
11
+ - A **TensorRT engine** is TensorRT's executable form of a model, optimized for a GPU, precision, and TensorRT version. It can be held in memory or saved to disk.
12
+ - **Building an engine** means creating that executable form from the model package's C++ network-building function and safetensors weights.
13
+ - A **cubin plugin** is a native TensorRT plugin whose GPU kernels are supplied as compiled CUDA binary (`.cubin`) files. The model package owns its CUDA C++ kernels and plugin behavior.
14
+
15
+ | Role | Use | Additional requirements |
16
+ | --- | --- | --- |
17
+ | Inference application | Call the model package's inference binding. | NVIDIA GPU and driver; the model package manages CUDA buffers and streams. |
18
+ | Model package developer | Open and execute engines through tensoracle's C++ SDK, exposing inference through your own binding. | TensorRT 11 headers, CMake 3.26 or later, C++17 compiler, and a binding framework when needed. |
19
+ | Cubin plugin developer | Compile CUDA C++ kernels and package a native plugin with the model. | `tensoracle[producer]`, PyPI CUDA headers and tools, TensorRT SDK, CMake, and a C++17 compiler. |
20
+
21
+ v2.0.0 supplies a Linux x86_64 `manylinux_2_28`, `cp312-abi3` wheel, including WSL2. The package currently declares Python 3.12 and 3.13 support. Target GPUs are sm86, sm89, and sm120 with TensorRT 11; fp8 requires sm89 or sm120. Windows x86_64 users must keep using tensoracle 1.x until a later v2 minor release adds a Windows wheel.
22
+
23
+ The wheel contains the C++ public headers, `libtensoracle`, CMake package configuration, and cubin plugin helper sources. TensorRT is resolved through a wheel dependency; the NVIDIA driver is not bundled. Public C++ interfaces and saved engine files remain compatible within a major version.
24
+
25
+ ## Build and reopen an engine in C++
26
+
27
+ Include `tensoracle/builder.hpp` and link `tensoracle::tensoracle` from `find_package(tensoracle 2 CONFIG REQUIRED)`. Locate the wheel's SDK with:
28
+
29
+ ```bash
30
+ python -c 'from importlib.metadata import distribution; print(distribution("tensoracle").locate_file("tensoracle/cpp"))'
31
+ ```
32
+
33
+ Pass that directory in `CMAKE_PREFIX_PATH`. The model package supplies a callback taking `nvinfer1::INetworkDefinition&` and `const tensoracle::weights_view&`. tensoracle reads declared safetensors and keeps their values alive through the build.
34
+
35
+ This excerpt follows the [separate consumer wheel](https://github.com/sync-dev-org/tensoracle/tree/main/examples/wheel_consumer). It assumes the example's `build_network` function, four-element float32 `bias.safetensors`, and caller-owned current CUDA context. The [user guide](https://sync-dev-org.github.io/tensoracle/#build-and-reopen-an-engine) shows the callback; the complete example supplies allocation, error handling, and nanobind integration.
36
+
37
+ ```cpp
38
+ #include <tensoracle/builder.hpp>
39
+
40
+ tensoracle::builder_options options;
41
+ options.name = "wheel-consumer-add";
42
+ options.graph_version = "1";
43
+ options.weights = {"bias.safetensors"};
44
+ options.inputs = {{"input", {4}, {4}, {4}}};
45
+ options.io = {{"input", TENSORACLE_FLOAT32, {4}, true},
46
+ {"output", TENSORACLE_FLOAT32, {4}, false}};
47
+ options.workspace_limit = 64 * 1024 * 1024;
48
+ options.optimization_level = 0;
49
+ auto engine = tensoracle::open_engine(options, build_network);
50
+ ```
51
+
52
+ On a valid cache hit, `tensoracle::open_engine` skips the callback, including after a process restart. Engine identity includes weights, plugin bytes, name, graph version, input profiles, builder optimization level, GPU compatibility properties, OS, CPU architecture, and TensorRT version. Update the graph version when network code changes; use separate names or versions for different precisions. Tensor types in the network determine precision, and TF32 is disabled.
53
+
54
+ Leave `workspace_limit` unset to use TensorRT's default or supply a positive byte count. It does not enter the cache identity; set `options.force_rebuild = true` to rebuild with a changed limit. Saved files live under `XDG_CACHE_HOME/tensoracle`, or `~/.cache/tensoracle` when unset. The [user guide](https://sync-dev-org.github.io/tensoracle/#save-and-reuse-an-engine) describes integrity checks, errors, and weight digest reuse.
55
+
56
+ ## Execute through the model package
57
+
58
+ The C++ model package calls `engine.create_execution()`, sets ranged input shapes together, queries output shapes and types, and enqueues every I/O device address on its own CUDA stream. Each execution unit retains the engine and owns an independent TensorRT context. The caller owns allocation, transfers, stream ordering, and completion checks. tensoracle does not synchronize enqueue or release; wait before changing shapes, freeing buffers, or closing the unit.
59
+
60
+ The [execution guide](https://sync-dev-org.github.io/tensoracle/#execute-an-engine) follows the consumer wheel's C++ API usage. The model package exposes its inference function to Python. Python applications do not construct or execute engines through tensoracle's Python utilities.
61
+
62
+ For CUDA Graphs, warm up one enqueue after each shape change before capturing on a nondefault stream. The graph fixes shapes, addresses, and execution unit; keep resources alive until work completes and the graph is destroyed. The caller creates and recreates graphs. Engines with loops, conditionals, or host-synchronizing layers or plugins cannot be captured. See the [capture guide](https://sync-dev-org.github.io/tensoracle/#capture-and-replay-with-cuda-graphs).
63
+
64
+ ## Manage saved engines
65
+
66
+ ```bash
67
+ tensoracle cache list
68
+ tensoracle cache prune --unused-for 30 --dry-run
69
+ tensoracle cache prune --unused-for 30
70
+ tensoracle cache prune --max-size 10GiB --dry-run
71
+ ```
72
+
73
+ These commands use the same cache directory as the C++ engine opener. They inspect headers and metadata without a GPU, original weights, or plan deserialization. Listing reports conditions, size, recorded time, workspace limit, and known weight paths. Empty files are `building or leftover`; invalid headers, including 1.x engine files, are `unrecognized`. Pruning never deletes either category.
74
+
75
+ An engine file's modification time records its initial save or last successful open. It does not indicate whether inference continues. Failed loads and management commands do not change it. Timestamp write failures do not prevent a successful load; records that could not be written, or predate timestamp recording, may be stale.
76
+
77
+ Pruning requires exactly one of `--unused-for DAYS` and `--max-size SIZE`. It prints a plan and rebuild notice before deletion and separates removed, excluded, missing, and failed files in its result. `--dry-run` changes nothing. Run pruning during maintenance without concurrent builds or opens. The [management guide](https://sync-dev-org.github.io/tensoracle/#manage-saved-engines) describes sizes, selection, and exit statuses.
78
+
79
+ ## Build a cubin plugin
80
+
81
+ Install `tensoracle[producer]` to generate sm86, sm89, and sm120 cubins from CUDA C++ source. The C++ producer uses NVRTC, `ptxas >=13.2,<14`, and `cuobjdump` from PyPI CUDA 13.2 wheels, without a GPU or OS-installed CUDA Toolkit. Triton kernels are not accepted. See the [producer API](https://sync-dev-org.github.io/tensoracle/api/producer/) for `KernelSpec` and `compile_bundle`.
82
+
83
+ Build the native plugin with [the wheel's C++ sources and CMake helpers](https://sync-dev-org.github.io/tensoracle/cpp-python-repositories/) and a separately prepared [TensorRT SDK](https://sync-dev-org.github.io/tensoracle/tensorrt-sdk/). Export TensorRT's `getCreators` and list all required libraries in `tensoracle::builder_options::plugins`. tensoracle inspects and registers them before construction or deserialization and hashes their bytes into the engine identity. Explicit registration is also available through `tensoracle/plugins.hpp`.
84
+
85
+ Package the plugin library and cubin bundle in your model wheel. `tensoracle.bundle.inspect_bundle` and `inspect_wheel` verify the manifest, cubin digests and actual architectures, declared argument ABI, and the plugin's embedded cubins without loading it. The manifest records NVRTC and ptxas versions and kernel source digests. The [plugin guide](https://sync-dev-org.github.io/tensoracle/#build-a-cubin-plugin) describes tensor types, formats, and the common lifecycle helper.
86
+
87
+ Follow the [four-stage optimization workflow](https://sync-dev-org.github.io/tensoracle/optimization/) and [precision guide](https://sync-dev-org.github.io/tensoracle/precision/) when validating a model or plugin.
88
+
89
+ ## Convert a PyTorch checkpoint
90
+
91
+ Install the conversion extra in the environment where you prepare weights:
92
+
93
+ ```bash
94
+ pip install 'tensoracle[convert]'
95
+ python -m tensoracle.convert model.pth model.safetensors
96
+ ```
97
+
98
+ Build scripts can call the same converter:
99
+
100
+ ```python
101
+ from tensoracle.convert import pth_to_safetensors
102
+
103
+ pth_to_safetensors("model.pth", "model.safetensors")
104
+ ```
105
+
106
+ The input must be a mapping with string keys whose values are tensors or nested mappings. Nested keys are joined with `.`: `{"state_dict": {"layer.weight": tensor}}` becomes `state_dict.layer.weight`. The converter preserves each tensor's dtype, shape, and values, including shared or noncontiguous dense tensors. It rejects non-tensor metadata and duplicate flattened names before writing the destination. Weight renaming and model-specific preprocessing belong to the model package.
107
+
108
+ The converter loads weights on the CPU with `torch.load(..., weights_only=True)` and rejects checkpoint objects instead of falling back to unrestricted pickle loading. The extra requires torch 2.6 or later, which includes the fix for [CVE-2025-32434](https://github.com/pytorch/pytorch/security/advisories/GHSA-53q9-r3pm-6pq6). Torch is imported only when conversion is called; it is absent from the package's required dependencies and is not imported by engine building or reuse. Calling the converter without torch raises `ImportError` with the installation command for `tensoracle[convert]`.
109
+
110
+ ## Compare output arrays
111
+
112
+ Run the TensorRT engines in your application, then pass one CPU NumPy output array and its corresponding reference array to `compare_outputs`. The first argument is the candidate; the second is the reference. Both arrays must have the same nonempty shape and be floating point, integer, or boolean arrays of the same kind. The function does not run an engine or decide whether its output passes a quality limit.
113
+
114
+ For floating point arrays, `OutputComparison` contains the non-finite count for each side, the maximum absolute error, the relative L2 error (difference norm divided by reference norm), and cosine similarity. Supply both `atol` and `rtol` to also get the count and fraction of elements outside `abs(candidate - reference) <= atol + rtol * abs(reference)`. Choose these limits for each output in your application. Floating point calculations use float64, including comparisons between fp16 and fp32 arrays. A non-finite value on either side makes all floating point measurements `None`. A zero reference norm makes relative L2 error `None`; a zero norm on either side makes cosine similarity `None`.
115
+
116
+ First, compare a saved correct output with the output from an fp32 TensorRT engine. The application creates and fixes the correct output outside tensoracle and runs the engine itself:
117
+
118
+ ```python
119
+ import numpy as np
120
+
121
+ from tensoracle import compare_outputs
122
+
123
+ correct_output = np.array([1.0, 2.0], dtype=np.float32) # Fixed reference output
124
+ fp32_output = np.array([1.0001, 2.0002], dtype=np.float32) # Engine output on CPU
125
+ fp32_result = compare_outputs(fp32_output, correct_output, atol=0.001, rtol=0.0)
126
+ fp32_within_tolerance = (
127
+ fp32_result.candidate_nonfinite_count == 0
128
+ and fp32_result.reference_nonfinite_count == 0
129
+ and fp32_result.outside_tolerance_count == 0
130
+ )
131
+ ```
132
+
133
+ Then use the fp32 TensorRT engine output as the reference for an fp16 TensorRT engine output. The application chooses limits for whole-array closeness or direction as well as any elementwise tolerance:
134
+
135
+ ```python
136
+ fp16_output = np.array([1.0, 2.0], dtype=np.float16) # Engine output on CPU
137
+ fp16_result = compare_outputs(fp16_output, fp32_output, atol=0.001, rtol=0.001)
138
+ fp16_accepted = (
139
+ fp16_result.candidate_nonfinite_count == 0
140
+ and fp16_result.reference_nonfinite_count == 0
141
+ and fp16_result.outside_tolerance_count == 0
142
+ and fp16_result.relative_l2_error is not None
143
+ and fp16_result.relative_l2_error <= 0.01 # Application's limit
144
+ )
145
+ ```
146
+
147
+ When tolerances are omitted, `outside_tolerance_count` and `outside_tolerance_fraction` are `None`. For integer and boolean arrays, use `mismatch_count` and `mismatch_fraction`; these measure exact equality instead of floating point distances. A zero mismatch count means those arrays agree. Shape, kind, or empty-array errors raise `ValueError`.
@@ -0,0 +1,213 @@
1
+ #include <nanobind/nanobind.h>
2
+ #include <nanobind/stl/string.h>
3
+ #include <tensoracle/cache.hpp>
4
+ #include <tensoracle/tensoracle.hpp>
5
+
6
+ #include <vector>
7
+
8
+ namespace nb = nanobind;
9
+
10
+ void bind_cubin(nb::module_& module);
11
+ namespace {
12
+ tensoracle_dtype array_dtype(nb::handle dtype) {
13
+ const auto kind = nb::cast<std::string>(dtype.attr("kind"));
14
+ const auto size = nb::cast<size_t>(dtype.attr("itemsize"));
15
+ if (kind == "b" && size == 1) return TENSORACLE_BOOL;
16
+ if (kind == "i" || kind == "u") {
17
+ const bool signed_type = kind == "i";
18
+ switch (size) {
19
+ case 1: return signed_type ? TENSORACLE_INT8 : TENSORACLE_UINT8;
20
+ case 2: return signed_type ? TENSORACLE_INT16 : TENSORACLE_UINT16;
21
+ case 4: return signed_type ? TENSORACLE_INT32 : TENSORACLE_UINT32;
22
+ case 8: return signed_type ? TENSORACLE_INT64 : TENSORACLE_UINT64;
23
+ }
24
+ }
25
+ if (kind == "f") {
26
+ if (size == 2) return TENSORACLE_FLOAT16;
27
+ if (size == 4) return TENSORACLE_FLOAT32;
28
+ if (size == 8) return TENSORACLE_FLOAT64;
29
+ if (size == sizeof(long double)) return TENSORACLE_LONG_DOUBLE;
30
+ }
31
+ return 0; // The shared library rejects unsupported numeric kinds.
32
+ }
33
+
34
+ struct cpu_array {
35
+ nb::object storage;
36
+ std::vector<size_t> shape;
37
+ Py_buffer buffer{};
38
+ tensoracle_array value{};
39
+
40
+ cpu_array(nb::handle numpy, nb::handle input) {
41
+ auto dtype = input.attr("dtype");
42
+ const auto type = array_dtype(dtype);
43
+ for (auto dimension : nb::borrow<nb::tuple>(input.attr("shape")))
44
+ shape.push_back(nb::cast<size_t>(dimension));
45
+ storage = numpy.attr("ascontiguousarray")(
46
+ input, nb::arg("dtype") = dtype.attr("newbyteorder")("="));
47
+ if (PyObject_GetBuffer(storage.ptr(), &buffer, PyBUF_SIMPLE) != 0)
48
+ throw nb::python_error();
49
+ value = {sizeof(value), type, shape.size(), shape.data(), buffer.buf};
50
+ }
51
+ cpu_array(const cpu_array&) = delete;
52
+ cpu_array& operator=(const cpu_array&) = delete;
53
+ ~cpu_array() { if (buffer.obj) PyBuffer_Release(&buffer); }
54
+ };
55
+
56
+ double tolerance_value(nb::handle input) {
57
+ try {
58
+ return static_cast<double>(nb::float_(input));
59
+ } catch (const nb::python_error& error) {
60
+ if (error.matches(PyExc_TypeError) || error.matches(PyExc_ValueError))
61
+ throw nb::value_error("atol and rtol must be finite nonnegative numbers");
62
+ throw;
63
+ }
64
+ }
65
+
66
+ template <class T>
67
+ nb::object optional_value(const tensoracle_output_comparison& result, uint32_t flag, T value) {
68
+ return result.available & flag ? nb::cast(value) : nb::none();
69
+ }
70
+
71
+ nb::tuple compare_outputs(nb::handle candidate, nb::handle reference, nb::handle atol, nb::handle rtol) {
72
+ auto numpy = nb::module_::import_("numpy");
73
+ for (auto input : {candidate, reference}) {
74
+ const auto is_array = PyObject_IsInstance(input.ptr(), numpy.attr("ndarray").ptr());
75
+ if (is_array < 0) throw nb::python_error();
76
+ if (!is_array) throw nb::type_error("candidate and reference must be NumPy arrays");
77
+ }
78
+ cpu_array a(numpy, candidate), b(numpy, reference);
79
+ tensoracle_comparison_options options{sizeof(options), !atol.is_none(), !rtol.is_none(), 0, 0};
80
+ // Presence and numeric validation belong to C++; these calls only convert
81
+ // Python values when both have been supplied.
82
+ if (options.has_atol && options.has_rtol) {
83
+ options.atol = tolerance_value(atol);
84
+ options.rtol = tolerance_value(rtol);
85
+ }
86
+ tensoracle_output_comparison result{};
87
+ result.struct_size = sizeof(result);
88
+ tensoracle_error* diagnostic = nullptr;
89
+ tensoracle_status status;
90
+ {
91
+ nb::gil_scoped_release release;
92
+ status = tensoracle_compare_outputs(&a.value, &b.value, &options, &result, &diagnostic);
93
+ }
94
+ try {
95
+ tensoracle::check(status, diagnostic);
96
+ } catch (const tensoracle::error& error) {
97
+ if (error.status() == TENSORACLE_INVALID_ARGUMENT) throw nb::value_error(error.what());
98
+ throw;
99
+ }
100
+ return nb::make_tuple(
101
+ result.candidate_nonfinite_count, result.reference_nonfinite_count,
102
+ optional_value(result, TENSORACLE_HAS_MAX_ABSOLUTE_ERROR, result.max_absolute_error),
103
+ optional_value(result, TENSORACLE_HAS_RELATIVE_L2_ERROR, result.relative_l2_error),
104
+ optional_value(result, TENSORACLE_HAS_COSINE_SIMILARITY, result.cosine_similarity),
105
+ optional_value(result, TENSORACLE_HAS_TOLERANCE, result.outside_tolerance_count),
106
+ optional_value(result, TENSORACLE_HAS_TOLERANCE, result.outside_tolerance_fraction),
107
+ optional_value(result, TENSORACLE_HAS_MISMATCH, result.mismatch_count),
108
+ optional_value(result, TENSORACLE_HAS_MISMATCH, result.mismatch_fraction));
109
+ }
110
+
111
+ template <class F> auto cache_operation(F&& operation) -> decltype(operation()) {
112
+ try {
113
+ nb::gil_scoped_release release;
114
+ return operation();
115
+ } catch (const tensoracle::error& error) {
116
+ if (error.status() == TENSORACLE_INVALID_ARGUMENT) throw nb::value_error(error.what());
117
+ throw;
118
+ }
119
+ }
120
+
121
+ int64_t cache_integer(nb::handle input) {
122
+ auto value = PyLong_AsLongLong(input.ptr());
123
+ if (PyErr_Occurred()) {
124
+ nb::python_error error;
125
+ if (error.matches(PyExc_OverflowError))
126
+ throw nb::value_error("cache prune criterion must fit in a signed 64-bit integer");
127
+ throw error;
128
+ }
129
+ return static_cast<int64_t>(value);
130
+ }
131
+
132
+ nb::dict cache_summary(const tensoracle::cache_result& owner) {
133
+ const auto summary = owner.summary();
134
+ nb::dict result;
135
+ result["directory"] = nb::str(summary.directory);
136
+ result["phase"] = summary.phase;
137
+ result["dry_run"] = bool(summary.dry_run);
138
+ result["total_size_bytes"] = summary.total_size_bytes;
139
+ result["removed_size_bytes"] = summary.removed_size_bytes;
140
+ result["remaining_size_bytes"] = summary.remaining_size_bytes;
141
+ result["projected_size_bytes"] = summary.projected_size_bytes;
142
+ result["max_size_bytes"] = summary.has_max_size ? nb::cast(summary.max_size_bytes) : nb::none();
143
+ result["max_size_reached"] = bool(summary.max_size_reached);
144
+ result["sizes_complete"] = bool(summary.sizes_complete);
145
+ result["succeeded"] = bool(summary.succeeded);
146
+ return result;
147
+ }
148
+
149
+ nb::list cache_entries(const tensoracle::cache_result& owner, const std::string& collection) {
150
+ int32_t kind;
151
+ if (collection == "entries") kind = TENSORACLE_CACHE_ENTRIES;
152
+ else if (collection == "targets") kind = TENSORACLE_CACHE_TARGETS;
153
+ else if (collection == "excluded") kind = TENSORACLE_CACHE_EXCLUDED;
154
+ else if (collection == "removed") kind = TENSORACLE_CACHE_REMOVED;
155
+ else if (collection == "failed") kind = TENSORACLE_CACHE_FAILED;
156
+ else if (collection == "missing") kind = TENSORACLE_CACHE_MISSING;
157
+ else throw nb::value_error("unknown cache result collection");
158
+ auto loads = nb::module_::import_("json").attr("loads");
159
+ nb::list output;
160
+ for (const auto& entry : owner.entries(kind)) {
161
+ nb::dict result;
162
+ result["file"] = nb::str(entry.file);
163
+ result["status"] = entry.status == TENSORACLE_CACHE_RECOGNIZED ? "recognized"
164
+ : entry.status == TENSORACLE_CACHE_BUILDING_OR_LEFTOVER ? "building or leftover"
165
+ : "unrecognized";
166
+ result["size_bytes"] = entry.has_stat ? nb::cast(entry.size_bytes) : nb::none();
167
+ result["recorded_at_seconds"] = entry.has_stat ? nb::cast(entry.recorded_at_seconds) : nb::none();
168
+ result["recorded_at_nanoseconds"] = entry.has_stat ? nb::cast(entry.recorded_at_nanoseconds) : nb::none();
169
+ // All C strings borrow owner. Copy/parse them while its handle is alive.
170
+ if (entry.header_json) result["header"] = loads(nb::str(entry.header_json));
171
+ if (entry.weight_paths_json) result["weight_paths"] = loads(nb::str(entry.weight_paths_json));
172
+ if (entry.reason) result["reason"] = nb::str(entry.reason);
173
+ output.append(result);
174
+ }
175
+ return output;
176
+ }
177
+ }
178
+
179
+ NB_MODULE(_native, module) {
180
+ bind_cubin(module);
181
+ module.doc() = "Thin bindings to tensoracle's shared C ABI";
182
+ module.def("version", [] {
183
+ auto version = tensoracle::version();
184
+ return nb::make_tuple(version.major, version.minor, version.patch);
185
+ });
186
+ module.def("tensorrt_version", [] {
187
+ auto version = tensoracle::tensorrt_version();
188
+ return nb::make_tuple(version.major, version.minor, version.patch);
189
+ });
190
+ module.def("compare_outputs", &compare_outputs, nb::arg("candidate"), nb::arg("reference"),
191
+ nb::kw_only(), nb::arg("atol") = nb::none(), nb::arg("rtol") = nb::none());
192
+ nb::class_<tensoracle::cache_result>(module, "_CacheResult")
193
+ .def("summary", &cache_summary)
194
+ .def("entries", &cache_entries, nb::arg("collection") = "entries");
195
+ module.def("cache_list", [] {
196
+ return cache_operation([] { return tensoracle::list_cache(); });
197
+ });
198
+ module.def("cache_plan_prune", [](nb::object unused_for, nb::object max_size) {
199
+ tensoracle_cache_prune_options options{sizeof(options), 0, 0, 0};
200
+ if (!unused_for.is_none()) {
201
+ options.selection |= TENSORACLE_CACHE_UNUSED_FOR;
202
+ options.unused_for_ns = cache_integer(unused_for);
203
+ }
204
+ if (!max_size.is_none()) {
205
+ options.selection |= TENSORACLE_CACHE_MAX_SIZE;
206
+ options.max_size_bytes = cache_integer(max_size);
207
+ }
208
+ return cache_operation([&] { return tensoracle::plan_cache_prune(options); });
209
+ }, nb::kw_only(), nb::arg("unused_for_ns") = nb::none(), nb::arg("max_size_bytes") = nb::none());
210
+ module.def("cache_apply_prune", [](const tensoracle::cache_result& plan, bool dry_run) {
211
+ return cache_operation([&] { return tensoracle::apply_cache_prune(plan, dry_run); });
212
+ }, nb::arg("plan"), nb::kw_only(), nb::arg("dry_run") = false);
213
+ }