MCPcopy
hub / github.com/NVIDIA/TensorRT-LLM / get_output_shapes

Method get_output_shapes

tensorrt_llm/python_plugin.py:363–375  ·  view source on GitHub ↗
(self, inputs, shape_inputs, exprBuilder)

Source from the content-addressed store, hash-verified

361 return [i.dtype for i in ret]
362
363 def get_output_shapes(self, inputs, shape_inputs, exprBuilder):
364 assert len(shape_inputs) == 0, "Currently we do not support shape inputs"
365
366 ret = self.shape_dtype_inference(
367 [SymTensor(None, ShapeExpr(i, exprBuilder)) for i in inputs]
368 )
369
370 ret = _convert_return_value_to_list(ret)
371 assert len(ret) == self.num_outputs
372 for i in ret:
373 assert isinstance(i, SymTensor)
374
375 return [i.shape.to_trt() for i in ret]
376
377 def supports_format_combination(self, pos, in_out, num_inputs):
378 """By default, TRT-LLM plugin supports all dtype and linear format.

Callers

nothing calls this directly

Calls 5

shape_dtype_inferenceMethod · 0.95
SymTensorClass · 0.85
ShapeExprClass · 0.85
to_trtMethod · 0.80

Tested by

no test coverage detected