tensoracle 0.3.1__tar.gz → 0.4.0__tar.gz

{tensoracle-0.3.1 → tensoracle-0.4.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.5
 Name: tensoracle
-Version: 0.3.1
+Version: 0.4.0
 Summary: Build, save, reuse, and execute TensorRT engines from safetensors weights
 Project-URL: documentation, https://sync-dev-org.github.io/tensoracle/
 Project-URL: repository, https://github.com/sync-dev-org/tensoracle
@@ -138,7 +138,7 @@ handle = open_engine(
 engine = handle.engine  # Keep handle.runtime alive while using the engine.
 ```
 
-Saved files live in the `tensoracle` directory under `XDG_CACHE_HOME`, or under `~/.cache` when that variable is unset. The cache identity includes the contents of all declared safetensors weight files and optional `plugins` files, the engine name, graph version, precision, input shapes, builder settings, GPU product name, OS, CPU architecture, and TensorRT version. Raise `graph_version` when graph code changes; without a version change, the previous engine is reused. Pass `force_rebuild=True` to replace a cached engine even when it is intact.
+Saved files live in the `tensoracle` directory under `XDG_CACHE_HOME`, or under `~/.cache` when that variable is unset. The cache identity includes the contents of all declared safetensors weight files and optional `plugins` files, the engine name, graph version, precision, input shapes, builder settings, GPU product name, OS, CPU architecture, and TensorRT version. Raise `graph_version` when graph code changes; without a version change, the previous engine is reused. Include in `graph_version` any other caller value that changes the engine contents but is not part of the cache identity, such as a plugin manifest or a build recipe. Pass `force_rebuild=True` to replace a cached engine even when it is intact.
 
 `plugins` names native TensorRT plugin libraries. Both `open_engine` and `build_serialized_engine` accept it. Each function checks that every library exports `getCreators`, loads it through TensorRT's registry, and retains it for the process before building the graph or deserializing an engine. Repeating the same library does not register it twice. `open_engine` also hashes the loaded library bytes into the engine identity; a changed library selects a different saved engine in a new process. A missing, unreadable, or incompatible library, or a library changed after loading in the same process, raises `PluginLibraryError` before a cached plan is opened. Keep the library and its cubin bundle together in the caller's wheel.
 
@@ -253,6 +253,8 @@ The caller selects the CUDA device and context, orders inputs written on other s
 
 Tensoracle neither allocates or frees device buffers, copies data, nor calls stream or device synchronization. TensorRT itself may allocate internal resources or act synchronously for some networks and plugins. Execution units reject HOST-location tensors, externally supplied shape tensor values, and data-dependent output sizes at creation. Ordinary networks that compute shapes internally from input dimensions are supported. CUDA graphs, profile switching, and caller-owned TensorRT workspace memory are outside this API.
 
+Networks with outputs marked as shape tensors emit `tensoracle.ShapeTensorOutputWarning` during building and when creating an execution unit (`create_execution()` or `EngineExecution(handle)`). The warning names the affected tensors: `create_execution()` cannot handle these engines, and calling it crashes the process. Building continues by default. To stop either operation safely, turn this warning into an exception with `warnings.filterwarnings("error", category=tensoracle.ShapeTensorOutputWarning)` after importing `warnings` and `tensoracle`. Saved engines retain the shape output names outside the cache conditions; older files without these names are reused without an execution warning.
+
 ## Build a cubin plugin
 
 Follow [Prepare a TensorRT SDK](https://github.com/sync-dev-org/tensoracle/blob/main/docs/tensorrt-sdk.md) to acquire and verify the headers and native library for a build SDK, including the fixed inputs for SDK 11.2.1.2 and the `TENSORRT_ROOT` setup.
@@ -267,6 +269,23 @@ pip install 'tensoracle[producer]'
 
 The wheel also contains `tensoracle/cpp/include`, `tensoracle/cpp/src`, and `tensoracle/cpp/cmake`. Pass the installed `tensoracle/cpp` directory to CMake as `TENSORACLE_CPP_ROOT` and include `TensoracleCubinPlugin.cmake`. Its `tensoracle_add_cubin_plugin` function compiles the common C++ source, the caller's plugin source, and the generated `bundle.cpp` into the caller's library. The caller implements the graph's meaning, field meanings, weight placement, launch sequence, and variant selection. The common source owns typed fields in their original order, handles `IPluginV3` lifecycle methods, checks cubin digests and resources, loads one module per CUDA context, checks the generated argument count and specialization rules, appends the two supported zero scratch pointers, and launches kernels. The caller selects a variant using `Kernel::accepts` and derives the count passed to `Kernel::launch` from its argument array. The common source releases modules and device allocations in their owning contexts. The library must export TensorRT's `getCreators`; a static registration constructor is not supported.
 
+`tensoracle::cubin::PluginBase` accepts a `PluginIO` declaration as its fourth constructor argument. `inputs` lists one `TensorSpec` per input in network order; each spec declares a TensorRT `type` and `format`. `output` declares the single output's type and format, and `shape_input_index` selects the input whose shape the output copies. For example, declare two fp16 / HWC8 inputs and one fp16 / HWC8 output with the shape of input 0:
+
+```cpp
+#include "tensoracle/cubin_plugin.hpp"
+
+tensoracle::cubin::PluginIO residual_io()
+{
+    using namespace tensoracle::cubin;
+    TensorSpec half{nvinfer1::DataType::kHALF, nvinfer1::TensorFormat::kHWC8};
+    return PluginIO{{half, half}, half, 0};
+}
+```
+
+Use `PluginBase(name, plugin_namespace, std::move(fields), residual_io())` in the derived plugin's constructor. For one fp16 / HWC8 input, use `PluginIO{{half}, half, 0}`. The base class checks tensor counts, types, and formats during format negotiation, configuration, and shape changes. Implement `validate_shape(nvinfer1::Dims const* inputs, std::int32_t input_count, nvinfer1::Dims const& output) const noexcept` to check the dimensions of every input in declaration order and the output; configuration may include wildcard dimensions (`-1`). For a residual operation, this callback can reject different shapes for its two inputs. The caller's kernel must handle the declared storage format. Omitting the fourth argument keeps one fp32 / LINEAR input and output with input 0's shape, so `PluginBase(name, plugin_namespace, fields)` remains valid. An empty input list or an invalid `shape_input_index` throws `std::invalid_argument` at construction.
+
+For configurations the declaration cannot express, such as multiple outputs or computed output shapes, developers can override the virtual V3 methods in a derived class or ask tensoracle's maintainers to extend the declaration.
+
 `tensoracle.bundle.inspect_bundle(bundle_directory, plugin=library_path)` checks the manifest, each cubin's SHA-256 and actual ELF SM, the supported hidden arguments and zero scratch sizes, and the cubin bytes embedded in the library. `inspect_wheel(wheel_path, plugin="path/inside/wheel.so")` performs the same checks on a wheel without installing it. Their distinct exceptions identify a missing cubin, corrupt cubin, false SM claim, or plugin and bundle mismatch. The manifest records Triton, ptxas, and kernel source digests without a producer host or OS path. The TensorRT versions tested with the included example are listed below.
 
 ## Tested environments

{tensoracle-0.3.1 → tensoracle-0.4.0}/README.md

@@ -108,7 +108,7 @@ handle = open_engine(
 engine = handle.engine  # Keep handle.runtime alive while using the engine.
 ```
 
-Saved files live in the `tensoracle` directory under `XDG_CACHE_HOME`, or under `~/.cache` when that variable is unset. The cache identity includes the contents of all declared safetensors weight files and optional `plugins` files, the engine name, graph version, precision, input shapes, builder settings, GPU product name, OS, CPU architecture, and TensorRT version. Raise `graph_version` when graph code changes; without a version change, the previous engine is reused. Pass `force_rebuild=True` to replace a cached engine even when it is intact.
+Saved files live in the `tensoracle` directory under `XDG_CACHE_HOME`, or under `~/.cache` when that variable is unset. The cache identity includes the contents of all declared safetensors weight files and optional `plugins` files, the engine name, graph version, precision, input shapes, builder settings, GPU product name, OS, CPU architecture, and TensorRT version. Raise `graph_version` when graph code changes; without a version change, the previous engine is reused. Include in `graph_version` any other caller value that changes the engine contents but is not part of the cache identity, such as a plugin manifest or a build recipe. Pass `force_rebuild=True` to replace a cached engine even when it is intact.
 
 `plugins` names native TensorRT plugin libraries. Both `open_engine` and `build_serialized_engine` accept it. Each function checks that every library exports `getCreators`, loads it through TensorRT's registry, and retains it for the process before building the graph or deserializing an engine. Repeating the same library does not register it twice. `open_engine` also hashes the loaded library bytes into the engine identity; a changed library selects a different saved engine in a new process. A missing, unreadable, or incompatible library, or a library changed after loading in the same process, raises `PluginLibraryError` before a cached plan is opened. Keep the library and its cubin bundle together in the caller's wheel.
 
@@ -223,6 +223,8 @@ The caller selects the CUDA device and context, orders inputs written on other s
 
 Tensoracle neither allocates or frees device buffers, copies data, nor calls stream or device synchronization. TensorRT itself may allocate internal resources or act synchronously for some networks and plugins. Execution units reject HOST-location tensors, externally supplied shape tensor values, and data-dependent output sizes at creation. Ordinary networks that compute shapes internally from input dimensions are supported. CUDA graphs, profile switching, and caller-owned TensorRT workspace memory are outside this API.
 
+Networks with outputs marked as shape tensors emit `tensoracle.ShapeTensorOutputWarning` during building and when creating an execution unit (`create_execution()` or `EngineExecution(handle)`). The warning names the affected tensors: `create_execution()` cannot handle these engines, and calling it crashes the process. Building continues by default. To stop either operation safely, turn this warning into an exception with `warnings.filterwarnings("error", category=tensoracle.ShapeTensorOutputWarning)` after importing `warnings` and `tensoracle`. Saved engines retain the shape output names outside the cache conditions; older files without these names are reused without an execution warning.
+
 ## Build a cubin plugin
 
 Follow [Prepare a TensorRT SDK](https://github.com/sync-dev-org/tensoracle/blob/main/docs/tensorrt-sdk.md) to acquire and verify the headers and native library for a build SDK, including the fixed inputs for SDK 11.2.1.2 and the `TENSORRT_ROOT` setup.
@@ -237,6 +239,23 @@ pip install 'tensoracle[producer]'
 
 The wheel also contains `tensoracle/cpp/include`, `tensoracle/cpp/src`, and `tensoracle/cpp/cmake`. Pass the installed `tensoracle/cpp` directory to CMake as `TENSORACLE_CPP_ROOT` and include `TensoracleCubinPlugin.cmake`. Its `tensoracle_add_cubin_plugin` function compiles the common C++ source, the caller's plugin source, and the generated `bundle.cpp` into the caller's library. The caller implements the graph's meaning, field meanings, weight placement, launch sequence, and variant selection. The common source owns typed fields in their original order, handles `IPluginV3` lifecycle methods, checks cubin digests and resources, loads one module per CUDA context, checks the generated argument count and specialization rules, appends the two supported zero scratch pointers, and launches kernels. The caller selects a variant using `Kernel::accepts` and derives the count passed to `Kernel::launch` from its argument array. The common source releases modules and device allocations in their owning contexts. The library must export TensorRT's `getCreators`; a static registration constructor is not supported.
 
+`tensoracle::cubin::PluginBase` accepts a `PluginIO` declaration as its fourth constructor argument. `inputs` lists one `TensorSpec` per input in network order; each spec declares a TensorRT `type` and `format`. `output` declares the single output's type and format, and `shape_input_index` selects the input whose shape the output copies. For example, declare two fp16 / HWC8 inputs and one fp16 / HWC8 output with the shape of input 0:
+
+```cpp
+#include "tensoracle/cubin_plugin.hpp"
+
+tensoracle::cubin::PluginIO residual_io()
+{
+    using namespace tensoracle::cubin;
+    TensorSpec half{nvinfer1::DataType::kHALF, nvinfer1::TensorFormat::kHWC8};
+    return PluginIO{{half, half}, half, 0};
+}
+```
+
+Use `PluginBase(name, plugin_namespace, std::move(fields), residual_io())` in the derived plugin's constructor. For one fp16 / HWC8 input, use `PluginIO{{half}, half, 0}`. The base class checks tensor counts, types, and formats during format negotiation, configuration, and shape changes. Implement `validate_shape(nvinfer1::Dims const* inputs, std::int32_t input_count, nvinfer1::Dims const& output) const noexcept` to check the dimensions of every input in declaration order and the output; configuration may include wildcard dimensions (`-1`). For a residual operation, this callback can reject different shapes for its two inputs. The caller's kernel must handle the declared storage format. Omitting the fourth argument keeps one fp32 / LINEAR input and output with input 0's shape, so `PluginBase(name, plugin_namespace, fields)` remains valid. An empty input list or an invalid `shape_input_index` throws `std::invalid_argument` at construction.
+
+For configurations the declaration cannot express, such as multiple outputs or computed output shapes, developers can override the virtual V3 methods in a derived class or ask tensoracle's maintainers to extend the declaration.
+
 `tensoracle.bundle.inspect_bundle(bundle_directory, plugin=library_path)` checks the manifest, each cubin's SHA-256 and actual ELF SM, the supported hidden arguments and zero scratch sizes, and the cubin bytes embedded in the library. `inspect_wheel(wheel_path, plugin="path/inside/wheel.so")` performs the same checks on a wheel without installing it. Their distinct exceptions identify a missing cubin, corrupt cubin, false SM claim, or plugin and bundle mismatch. The manifest records Triton, ptxas, and kernel source digests without a producer host or OS path. The TensorRT versions tested with the included example are listed below.
 
 ## Tested environments

{tensoracle-0.3.1 → tensoracle-0.4.0}/cpp/include/tensoracle/cubin_plugin.hpp

@@ -41,13 +41,31 @@ private:
     nvinfer1::PluginFieldCollection collection_{};
 };
 
-// Implements TensorRT V3 lifecycle and the format/serialization boilerplate.
+struct TensorSpec {
+    nvinfer1::DataType type{nvinfer1::DataType::kFLOAT};
+    nvinfer1::TensorFormat format{nvinfer1::TensorFormat::kLINEAR};
+};
+
+// Declares the input count, each tensor's type and format, and the input whose
+// shape is copied to the single output. For configurations this cannot express
+// (multiple outputs or computed output shapes), developers can override virtual
+// V3 methods or ask tensoracle's maintainers to extend the declaration.
+struct PluginIO {
+    std::vector<TensorSpec> inputs{TensorSpec{}};
+    TensorSpec output{};
+    std::int32_t shape_input_index{0};
+};
+
+// Implements TensorRT V3 lifecycle and format/serialization boilerplate using
+// a PluginIO declaration. For other configurations, developers can override
+// virtual V3 methods or ask tensoracle's maintainers to extend the declaration.
 class PluginBase : public nvinfer1::IPluginV3,
                    public nvinfer1::IPluginV3OneCore,
                    public nvinfer1::IPluginV3OneBuild,
                    public nvinfer1::IPluginV3OneRuntime {
 public:
-    PluginBase(char const* name, char const* plugin_namespace, Fields fields);
+    // Throws std::invalid_argument for no inputs or an out-of-range shape input.
+    PluginBase(char const* name, char const* plugin_namespace, Fields fields, PluginIO io = {});
     ~PluginBase() noexcept override = default;
 
     nvinfer1::IPluginCapability* getCapabilityInterface(nvinfer1::PluginCapabilityType type) noexcept override;
@@ -79,13 +97,16 @@ public:
 
 protected:
     Fields const& fields() const noexcept { return fields_; }
-    virtual std::int32_t validate_shape(nvinfer1::Dims const& input,
+    // Inputs are in declaration order; the array is valid only during this call.
+    // configurePlugin may provide wildcard (-1) dimensions.
+    virtual std::int32_t validate_shape(nvinfer1::Dims const* inputs, std::int32_t input_count,
         nvinfer1::Dims const& output) const noexcept = 0;
 
 private:
     char const* name_;
     char const* namespace_;
     Fields fields_;
+    PluginIO io_;
 };
 
 enum class ArgumentKind : std::uint8_t { kPointer, kInt32, kFloat32 };

{tensoracle-0.3.1 → tensoracle-0.4.0}/cpp/src/cubin_plugin.cpp

@@ -4,11 +4,48 @@
 #include <array>
 #include <cstring>
 #include <limits>
+#include <stdexcept>
 #include <utility>
 
 namespace tensoracle::cubin {
 namespace {
 
+nvinfer1::PluginTensorDesc const& tensor_desc(nvinfer1::PluginTensorDesc const& value) noexcept
+{
+    return value;
+}
+
+nvinfer1::PluginTensorDesc const& tensor_desc(nvinfer1::DynamicPluginTensorDesc const& value) noexcept
+{
+    return value.desc;
+}
+
+bool matches_tensor(nvinfer1::PluginTensorDesc const& descriptor, TensorSpec const& spec) noexcept
+{
+    return descriptor.type == spec.type && descriptor.format == spec.format;
+}
+
+template <typename Descriptor, typename Validate>
+std::int32_t validate_tensors(Descriptor const* inputs, std::int32_t input_count,
+    Descriptor const* outputs, std::int32_t output_count, PluginIO const& io,
+    Validate validate) noexcept
+{
+    if (inputs == nullptr || outputs == nullptr || input_count < 0 || output_count != 1 ||
+        static_cast<std::size_t>(input_count) != io.inputs.size()) return -1;
+    for (std::int32_t index = 0; index < input_count; ++index)
+        if (!matches_tensor(tensor_desc(inputs[index]), io.inputs[index])) return -1;
+    if (!matches_tensor(tensor_desc(outputs[0]), io.output)) return -1;
+    try {
+        std::vector<nvinfer1::Dims> dims;
+        dims.reserve(io.inputs.size());
+        for (std::int32_t index = 0; index < input_count; ++index)
+            dims.push_back(tensor_desc(inputs[index]).dims);
+        return validate(dims.data(), input_count, tensor_desc(outputs[0]).dims);
+    } catch (...) {
+        return -1;
+    }
+}
+
 std::size_t field_size(nvinfer1::PluginFieldType type) noexcept
 {
     using nvinfer1::PluginFieldType;
@@ -102,8 +139,14 @@ void Fields::rebuild()
 
 nvinfer1::PluginFieldCollection const* Fields::collection() noexcept { return &collection_; }
 
-PluginBase::PluginBase(char const* name, char const* plugin_namespace, Fields fields)
-    : name_(name), namespace_(plugin_namespace), fields_(std::move(fields)) {}
+PluginBase::PluginBase(char const* name, char const* plugin_namespace, Fields fields, PluginIO io)
+    : name_(name), namespace_(plugin_namespace), fields_(std::move(fields)), io_(std::move(io))
+{
+    if (io_.inputs.empty() || io_.inputs.size() > static_cast<std::size_t>(
+            std::numeric_limits<std::int32_t>::max()) || io_.shape_input_index < 0 ||
+        static_cast<std::size_t>(io_.shape_input_index) >= io_.inputs.size())
+        throw std::invalid_argument("PluginIO requires inputs and a valid shape_input_index");
+}
 
 nvinfer1::IPluginCapability* PluginBase::getCapabilityInterface(
     nvinfer1::PluginCapabilityType type) noexcept
@@ -124,20 +167,23 @@ bool PluginBase::supportsFormatCombination(std::int32_t position,
     nvinfer1::DynamicPluginTensorDesc const* in_out,
     std::int32_t input_count, std::int32_t output_count) noexcept
 {
-    if (in_out == nullptr || input_count != 1 || output_count != 1 || position < 0 || position > 1)
+    if (in_out == nullptr || input_count < 0 ||
+        static_cast<std::size_t>(input_count) != io_.inputs.size() ||
+        output_count != 1 || position < 0 || position > input_count)
         return false;
-    auto const& descriptor = in_out[position].desc;
-    return descriptor.type == nvinfer1::DataType::kFLOAT &&
-        descriptor.format == nvinfer1::TensorFormat::kLINEAR;
+    auto const& spec = position < input_count ? io_.inputs[position] : io_.output;
+    return matches_tensor(in_out[position].desc, spec);
 }
 
 std::int32_t PluginBase::getOutputDataTypes(nvinfer1::DataType* output_types,
     std::int32_t output_count, nvinfer1::DataType const* input_types,
     std::int32_t input_count) const noexcept
 {
-    if (output_types == nullptr || input_types == nullptr || output_count != 1 || input_count != 1 ||
-        input_types[0] != nvinfer1::DataType::kFLOAT) return -1;
-    output_types[0] = nvinfer1::DataType::kFLOAT;
+    if (output_types == nullptr || input_types == nullptr || output_count != 1 || input_count < 0 ||
+        static_cast<std::size_t>(input_count) != io_.inputs.size()) return -1;
+    for (std::int32_t index = 0; index < input_count; ++index)
+        if (input_types[index] != io_.inputs[index].type) return -1;
+    output_types[0] = io_.output.type;
     return 0;
 }
 
@@ -146,9 +192,10 @@ std::int32_t PluginBase::getOutputShapes(nvinfer1::DimsExprs const* inputs,
     nvinfer1::DimsExprs* outputs, std::int32_t output_count,
     nvinfer1::IExprBuilder&) noexcept
 {
-    if (inputs == nullptr || outputs == nullptr || input_count != 1 ||
+    if (inputs == nullptr || outputs == nullptr || input_count < 0 ||
+        static_cast<std::size_t>(input_count) != io_.inputs.size() ||
         shape_input_count != 0 || output_count != 1) return -1;
-    outputs[0] = inputs[0];
+    outputs[0] = inputs[io_.shape_input_index];
     return 0;
 }
 
@@ -156,8 +203,10 @@ std::int32_t PluginBase::configurePlugin(nvinfer1::DynamicPluginTensorDesc const
     std::int32_t input_count, nvinfer1::DynamicPluginTensorDesc const* outputs,
     std::int32_t output_count) noexcept
 {
-    if (inputs == nullptr || outputs == nullptr || input_count != 1 || output_count != 1) return -1;
-    return validate_shape(inputs[0].desc.dims, outputs[0].desc.dims);
+    return validate_tensors(inputs, input_count, outputs, output_count, io_,
+        [this](nvinfer1::Dims const* dims, std::int32_t count, nvinfer1::Dims const& output) {
+            return validate_shape(dims, count, output);
+        });
 }
 
 std::size_t PluginBase::getWorkspaceSize(nvinfer1::DynamicPluginTensorDesc const*, std::int32_t,
@@ -169,8 +218,10 @@ std::int32_t PluginBase::onShapeChange(nvinfer1::PluginTensorDesc const* inputs,
     std::int32_t input_count, nvinfer1::PluginTensorDesc const* outputs,
     std::int32_t output_count) noexcept
 {
-    if (inputs == nullptr || outputs == nullptr || input_count != 1 || output_count != 1) return -1;
-    return validate_shape(inputs[0].dims, outputs[0].dims);
+    return validate_tensors(inputs, input_count, outputs, output_count, io_,
+        [this](nvinfer1::Dims const* dims, std::int32_t count, nvinfer1::Dims const& output) {
+            return validate_shape(dims, count, output);
+        });
 }
 
 nvinfer1::IPluginV3* PluginBase::attachToContext(nvinfer1::IPluginResourceContext*) noexcept

{tensoracle-0.3.1 → tensoracle-0.4.0}/docs/index.md

@@ -70,7 +70,7 @@ handle = open_engine(
 )
 ```
 
-`my_model` represents the application's model package. On a valid cache hit, `open_engine` does not call `graph`; on a miss, it builds and saves a new engine. `open_engine` itself does not execute the engine or manage GPU buffers. Create an execution unit from its returned handle to enqueue inference. Saved engines live under `XDG_CACHE_HOME/tensoracle`, or `~/.cache/tensoracle` when `XDG_CACHE_HOME` is unset. See the [complete cache example](https://github.com/sync-dev-org/tensoracle#save-and-reopen-an-engine) for the declaration and expected input and output types.
+`my_model` represents the application's model package. On a valid cache hit, `open_engine` does not call `graph`; on a miss, it builds and saves a new engine. Raise `graph_version` when the graph code changes or when any other caller value that changes the engine contents but is not part of the cache identity changes, such as a plugin manifest or a build recipe. `open_engine` itself does not execute the engine or manage GPU buffers. Create an execution unit from its returned handle to enqueue inference. Saved engines live under `XDG_CACHE_HOME/tensoracle`, or `~/.cache/tensoracle` when `XDG_CACHE_HOME` is unset. See the [complete cache example](https://github.com/sync-dev-org/tensoracle#save-and-reopen-an-engine) for the declaration and expected input and output types.
 
 ### Execute an engine
 
@@ -155,6 +155,8 @@ The caller selects the CUDA device and context, orders inputs written on other s
 
 Tensoracle neither allocates or frees device buffers, copies data, nor calls stream or device synchronization. TensorRT itself may allocate internal resources or act synchronously for some networks and plugins. Execution units reject HOST-location tensors, externally supplied shape tensor values, and data-dependent output sizes at creation. Ordinary networks that compute shapes internally from input dimensions are supported. CUDA graphs, profile switching, and caller-owned TensorRT workspace memory are outside this API.
 
+Networks with outputs marked as shape tensors emit `tensoracle.ShapeTensorOutputWarning` during building and when creating an execution unit (`create_execution()` or `EngineExecution(handle)`). The warning names the affected tensors: `create_execution()` cannot handle these engines, and calling it crashes the process. Building continues by default. To stop either operation safely, turn this warning into an exception with `warnings.filterwarnings("error", category=tensoracle.ShapeTensorOutputWarning)` after importing `warnings` and `tensoracle`. Saved engines retain the shape output names outside the cache conditions; older files without these names are reused without an execution warning.
+
 `TensorSpec` accepts optional `location=trt.TensorLocation.DEVICE` and `format=trt.TensorFormat.LINEAR`. `open_engine` checks these only when declared, together with tensor names, modes, dtypes, build-time shapes, and profiles. An `EngineIOError` names the tensor, field, observed value, expected value, and saved file path. The execution unit does not accept a second declaration.
 
 ### Build a standard graph
@@ -173,6 +175,23 @@ pip install 'tensoracle[producer]'
 
 `tensoracle.producer.compile_bundle` uses the extra's Triton 3.8.0, `ptxas` 12.9.86, and `cuobjdump` 13.3 to produce cubins for sm86, sm89, and sm120. It finds the NVIDIA wheel executables by default and accepts explicit paths as overrides. Cubin production does not need a GPU or import torch; run it without importing torch first. Build the native plugin with a TensorRT 11 SDK, CUDA toolkit, CMake, and a C++17 compiler. Inspect the resulting bundle or wheel with `tensoracle.bundle`. The [cubin plugin example](https://github.com/sync-dev-org/tensoracle/tree/main/examples/cubin_plugin) covers rebuilding from a repository clone, packaging, and runtime use. Normal import, engine building, and engine reuse do not import the producer or Triton, and do not need the producer extra or native build tools.
 
+Pass a `tensoracle::cubin::PluginIO` declaration as the fourth argument to `PluginBase(name, plugin_namespace, fields, io)`. The `inputs` vector determines the input count and declares each input's TensorRT `type` and `format` through `TensorSpec`; `output` declares the single output's type and format. `shape_input_index` selects the input whose shape the output copies. This example declares two fp16 / HWC8 inputs and an fp16 / HWC8 output with input 0's shape:
+
+```cpp
+#include "tensoracle/cubin_plugin.hpp"
+
+tensoracle::cubin::PluginIO residual_io()
+{
+    using namespace tensoracle::cubin;
+    TensorSpec half{nvinfer1::DataType::kHALF, nvinfer1::TensorFormat::kHWC8};
+    return PluginIO{{half, half}, half, 0};
+}
+```
+
+Use `PluginBase(name, plugin_namespace, std::move(fields), residual_io())` in the derived plugin's constructor; use `PluginIO{{half}, half, 0}` for one fp16 / HWC8 input. The common implementation negotiates the declared types and formats, sets the output type and shape, and checks counts, types, and formats during configuration and shape changes. Implement `validate_shape(nvinfer1::Dims const* inputs, std::int32_t input_count, nvinfer1::Dims const& output) const noexcept` to check all input dimensions in declaration order and the output dimensions, including wildcard dimensions (`-1`) during configuration. A residual plugin can reject mismatched input shapes in this callback. The kernel must handle the declared storage format. Omitting `io` keeps one fp32 / LINEAR input and output with input 0's shape; the original three-argument constructor call remains valid. An empty input list or an out-of-range `shape_input_index` throws `std::invalid_argument` at construction.
+
+For configurations the declaration cannot express, such as multiple outputs or computed output shapes, developers can override the virtual V3 methods in a derived class or ask tensoracle's maintainers to extend the declaration.
+
 ## API reference
 
 - [Engine building, cache reuse, execution, and output comparison](api/index.md)

{tensoracle-0.3.1 → tensoracle-0.4.0}/docs/optimization.md

@@ -18,7 +18,7 @@ Use the same representative inputs, shapes, and corresponding outputs throughout
 
 First check whether a combination of standard TensorRT layers can preserve the operation's meaning, shapes, control flow, and output quality. For example, a deformable convolution may be expressible through grid sampling, mask multiplication, and convolution. If a plugin bundled with TensorRT matches the required behavior, compare that option too. Use a custom cubin plugin only when the suitable alternatives still leave a bottleneck.
 
-- **tensoracle tools:** `tensoracle.producer.compile_bundle` generates cubins and a manifest for supported GPU generations. The shared C++ plugin helpers support a caller-built native plugin, and `tensoracle.bundle.inspect_bundle` or `inspect_wheel` checks the cubins and their match with the plugin. Pass the plugin library through `plugins=` to `build_serialized_engine` or `open_engine`. Follow [Build a cubin plugin](https://github.com/sync-dev-org/tensoracle#build-a-cubin-plugin) and the [cubin plugin example](https://github.com/sync-dev-org/tensoracle/blob/main/examples/cubin_plugin/README.md) for the generation, build, and inspection operations.
+- **tensoracle tools:** `tensoracle.producer.compile_bundle` generates cubins and a manifest for supported GPU generations. The shared C++ plugin helpers support a caller-built native plugin with declared tensor types, formats, and an input whose shape the output copies; `tensoracle.bundle.inspect_bundle` or `inspect_wheel` checks the cubins and their match with the plugin. Pass the plugin library through `plugins=` to `build_serialized_engine` or `open_engine`. Follow [Build a cubin plugin](index.md#build-a-cubin-plugin) for the `PluginIO` declaration, including fp16 / HWC8 and two-input examples, and the [cubin plugin example](https://github.com/sync-dev-org/tensoracle/blob/main/examples/cubin_plugin/README.md) for the generation, build, and inspection operations.
 - **You provide:** Identify the bottleneck, write the kernel and model-specific plugin behavior, choose its supported shapes and types, build and package its native library, and run the resulting engines. An fp32 kernel must not use TF32.
 - **Gate:** For a custom cubin plugin, the bundle and plugin inspections must pass. Reapply stage 1 to the replacement's fp32 engine against the same fixed PyTorch reference. If using fp16, reapply stage 2 against an accepted fp32 engine with the same limits and model-specific measurements. Advance only when the replacement meets those quality gates; otherwise retain the accepted implementation.
 

{tensoracle-0.3.1 → tensoracle-0.4.0}/examples/cubin_plugin/plugin.cpp

@@ -117,7 +117,7 @@ public:
     {
         if (!valid_ || input_descriptors == nullptr || output_descriptors == nullptr ||
             inputs == nullptr || outputs == nullptr || inputs[0] == nullptr || outputs[0] == nullptr ||
-            validate_shape(input_descriptors[0].dims, output_descriptors[0].dims) != 0) return -1;
+            validate_shape(&input_descriptors[0].dims, 1, output_descriptors[0].dims) != 0) return -1;
         int rows = input_descriptors[0].dims.d[0];
         if (rows < 1 || rows > INT_MAX / 16) return -1;
         int const index = sm_index();
@@ -141,9 +141,10 @@ public:
     }
 
 protected:
-    std::int32_t validate_shape(Dims const& input, Dims const& output) const noexcept override
+    std::int32_t validate_shape(Dims const* inputs, std::int32_t input_count,
+        Dims const& output) const noexcept override
     {
-        return shape_matches(input, output, 32) ? 0 : -1;
+        return input_count == 1 && shape_matches(inputs[0], output, 32) ? 0 : -1;
     }
 
 private:
@@ -178,7 +179,7 @@ public:
     {
         if (!valid_ || input_descriptors == nullptr || output_descriptors == nullptr ||
             inputs == nullptr || outputs == nullptr || inputs[0] == nullptr || outputs[0] == nullptr ||
-            validate_shape(input_descriptors[0].dims, output_descriptors[0].dims) != 0) return -1;
+            validate_shape(&input_descriptors[0].dims, 1, output_descriptors[0].dims) != 0) return -1;
         int const rows = input_descriptors[0].dims.d[0];
         if (rows < 1 || rows > INT_MAX / 16) return -1;
         int const index = sm_index();
@@ -194,9 +195,10 @@ public:
     }
 
 protected:
-    std::int32_t validate_shape(Dims const& input, Dims const& output) const noexcept override
+    std::int32_t validate_shape(Dims const* inputs, std::int32_t input_count,
+        Dims const& output) const noexcept override
     {
-        return shape_matches(input, output, 16) ? 0 : -1;
+        return input_count == 1 && shape_matches(inputs[0], output, 16) ? 0 : -1;
     }
 
 private:

{tensoracle-0.3.1 → tensoracle-0.4.0}/src/tensoracle/__init__.py

@@ -15,6 +15,7 @@ __all__ = [
     "InputShape",
     "OutputComparison",
     "PluginLibraryError",
+    "ShapeTensorOutputWarning",
     "TensorSpec",
     "build_serialized_engine",
     "compare_outputs",
@@ -23,7 +24,14 @@ __all__ = [
 
 _MODULES = {
     **dict.fromkeys(
-        ("BuilderSettings", "EngineBuild", "InputShape", "build_serialized_engine"), "._build"
+        (
+            "BuilderSettings",
+            "EngineBuild",
+            "InputShape",
+            "ShapeTensorOutputWarning",
+            "build_serialized_engine",
+        ),
+        "._build",
     ),
     **dict.fromkeys(
         (

{tensoracle-0.3.1 → tensoracle-0.4.0}/src/tensoracle/_build.py

@@ -1,5 +1,6 @@
 """Build serialized TensorRT engines from caller-defined network graphs."""
 
+import warnings
 from collections.abc import Callable, Mapping, Sequence
 from dataclasses import dataclass, field
 from pathlib import Path
@@ -11,6 +12,21 @@ from ._memory import release_host_memory
 from ._plugins import register_libraries
 
 
+class ShapeTensorOutputWarning(UserWarning):
+    """Shape tensor outputs cannot be handled by tensoracle execution units."""
+
+
+def _warn_shape_tensor_outputs(names: tuple[str, ...], *, stacklevel: int) -> None:
+    if names:
+        warnings.warn(
+            f"TensorRT outputs {', '.join(repr(name) for name in names)} are shape tensors; "
+            "tensoracle.create_execution() cannot handle this engine and calling it will "
+            "crash the process.",
+            ShapeTensorOutputWarning,
+            stacklevel=stacklevel + 1,
+        )
+
+
 @dataclass(frozen=True)
 class InputShape:
     """The minimum, preferred, and maximum dimensions of one input."""
@@ -76,6 +92,17 @@ def build_serialized_engine(
     logger: trt.ILogger | None = None,
 ) -> bytes:
     """Register ``plugins`` before calling ``graph`` and serializing its engine."""
+    plan, _ = _build_serialized_engine(graph, declaration, plugins=plugins, logger=logger)
+    return plan
+
+
+def _build_serialized_engine(
+    graph: Callable[[trt.INetworkDefinition], None],
+    declaration: EngineBuild,
+    *,
+    plugins: Sequence[Path] = (),
+    logger: trt.ILogger | None = None,
+) -> tuple[bytes, tuple[str, ...]]:
     if declaration.precision not in ("fp32", "fp16"):
         raise ValueError(f"unsupported precision: {declaration.precision}")
     if declaration.settings.tf32:
@@ -93,6 +120,13 @@ def build_serialized_engine(
         builder = trt.Builder(logger or trt.Logger(trt.Logger.ERROR))
         network = builder.create_network(0)
         graph(network)
+        shape_tensor_outputs = []
+        for index in range(network.num_outputs):
+            tensor = network.get_output(index)
+            if tensor.is_shape_tensor:
+                shape_tensor_outputs.append(tensor.name)
+        output_names = tuple(shape_tensor_outputs)
+        _warn_shape_tensor_outputs(output_names, stacklevel=3)
         config = builder.create_builder_config()
         config.clear_flag(trt.BuilderFlag.TF32)
         config.set_memory_pool_limit(
@@ -129,7 +163,7 @@ def build_serialized_engine(
         serialized = builder.build_serialized_network(network, config)
         if serialized is None:
             raise RuntimeError("TensorRT failed to build the serialized engine")
-        return bytes(serialized)
+        return bytes(serialized), output_names
     finally:
         # Tensor references also keep the network alive, including in exception frames.
         inputs.clear()

{tensoracle-0.3.1 → tensoracle-0.4.0}/src/tensoracle/_cache.py

@@ -9,12 +9,13 @@ import platform
 import tempfile
 from collections.abc import Callable, Mapping, Sequence
 from dataclasses import asdict, dataclass
+from dataclasses import field as dataclass_field
 from pathlib import Path
 from typing import TYPE_CHECKING
 
 import tensorrt as trt
 
-from ._build import EngineBuild, build_serialized_engine
+from ._build import EngineBuild, _build_serialized_engine
 from ._plugins import register_libraries
 
 if TYPE_CHECKING:
@@ -49,6 +50,7 @@ class EngineHandle:
     engine: trt.ICudaEngine
     runtime: trt.Runtime
     path: Path
+    _shape_tensor_outputs: tuple[str, ...] = dataclass_field(default=(), repr=False)
 
     def create_execution(self) -> "EngineExecution":
         """Create an independent execution context with explicit close semantics."""
@@ -103,16 +105,29 @@ def open_engine(
         try:
             fd = os.open(path, os.O_WRONLY | os.O_CREAT | os.O_EXCL, 0o600)
         except FileExistsError:
-            plan = _read_plan(path, conditions)
-            return _open_plan(plan, path, declaration, io, actual_logger)
+            plan, shape_tensor_outputs = _read_plan(path, conditions)
+            return _open_plan(
+                plan,
+                path,
+                declaration,
+                io,
+                actual_logger,
+                shape_tensor_outputs=shape_tensor_outputs,
+            )
         else:
             os.close(fd)
 
     temporary: Path | None = None
     try:
-        plan = build_serialized_engine(graph, declaration, plugins=plugins, logger=actual_logger)
-        handle = _open_plan(plan, path, declaration, io, actual_logger)
+        plan, shape_tensor_outputs = _build_serialized_engine(
+            graph, declaration, plugins=plugins, logger=actual_logger
+        )
+        handle = _open_plan(
+            plan, path, declaration, io, actual_logger, shape_tensor_outputs=shape_tensor_outputs
+        )
         header = {"conditions": conditions, "plan_sha256": hashlib.sha256(plan).hexdigest()}
+        if shape_tensor_outputs:
+            header["shape_tensor_outputs"] = list(shape_tensor_outputs)
         header_bytes = _json_bytes(header)
         fd, filename = tempfile.mkstemp(prefix=f".{digest}.", dir=path.parent)
         temporary = Path(filename)
@@ -279,7 +294,7 @@ def _conditions(
     }
 
 
-def _read_plan(path: Path, conditions: dict) -> memoryview:
+def _read_plan(path: Path, conditions: dict) -> tuple[memoryview, tuple[str, ...]]:
     data = path.read_bytes()
     offset = len(_MAGIC)
     if not data.startswith(_MAGIC) or len(data) < offset + 8:
@@ -294,12 +309,17 @@ def _read_plan(path: Path, conditions: dict) -> memoryview:
         raise CorruptEngineError(f"invalid TensorRT engine header: {path}") from error
     if not isinstance(header, dict) or not isinstance(header.get("conditions"), dict):
         raise CorruptEngineError(f"invalid TensorRT engine header: {path}")
+    shape_tensor_outputs = header.get("shape_tensor_outputs", [])
+    if not isinstance(shape_tensor_outputs, list) or any(
+        not isinstance(name, str) for name in shape_tensor_outputs
+    ):
+        raise CorruptEngineError(f"invalid TensorRT engine header: {path}")
     if header["conditions"] != conditions:
         raise EngineConditionError(f"TensorRT engine conditions differ: {path}")
     plan = memoryview(data)[header_end:]
     if header.get("plan_sha256") != hashlib.sha256(plan).hexdigest():
         raise CorruptEngineError(f"TensorRT engine checksum differs: {path}")
-    return plan
+    return plan, tuple(shape_tensor_outputs)
 
 
 def _open_plan(
@@ -308,6 +328,8 @@ def _open_plan(
     declaration: EngineBuild,
     io: EngineIO,
     logger: trt.ILogger,
+    *,
+    shape_tensor_outputs: tuple[str, ...] = (),
 ) -> EngineHandle:
     runtime = trt.Runtime(logger)
     engine = runtime.deserialize_cuda_engine(plan)
@@ -349,7 +371,7 @@ def _open_plan(
             actual_range = (io.inputs[name].shape,) * 3
         if actual_range != expected_range:
             _io_mismatch(path, name, "profile range", actual_range, expected_range)
-    return EngineHandle(engine, runtime, path)
+    return EngineHandle(engine, runtime, path, _shape_tensor_outputs=shape_tensor_outputs)
 
 
 def _io_mismatch(path: Path, name: str, field: str, observed: object, expected: object) -> None:

{tensoracle-0.3.1 → tensoracle-0.4.0}/src/tensoracle/_execution.py

@@ -7,6 +7,8 @@ from typing import TYPE_CHECKING, Self
 
 import tensorrt as trt
 
+from ._build import _warn_shape_tensor_outputs
+
 if TYPE_CHECKING:
     from ._cache import EngineHandle
 
@@ -51,6 +53,7 @@ class EngineExecution:
             self._context = engine.create_execution_context()
             if self._context is None:
                 raise RuntimeError("TensorRT failed to create an execution context")
+            _warn_shape_tensor_outputs(handle._shape_tensor_outputs, stacklevel=2)
             # Probe output dimensions before accepting the engine. A ranged engine
             # still requires an explicit caller setting after this capability check.
             self._set_shapes(initial_shapes, creating=True)

{tensoracle-0.3.1 → tensoracle-0.4.0}/src/tensoracle/_version.py

@@ -18,7 +18,7 @@ version_tuple: tuple[int | str, ...]
 commit_id: str | None
 __commit_id__: str | None
 
-__version__ = version = '0.3.1'
-__version_tuple__ = version_tuple = (0, 3, 1)
+__version__ = version = '0.4.0'
+__version_tuple__ = version_tuple = (0, 4, 0)
 
 __commit_id__ = commit_id = None