Method register_ndarray

tensorrt_llm/network.py:301–308

When the functional APIs need to create local numpy array and use as weights for constant or other layers, they need to register the ndarray objects to the TRT-LLM Network to prolong the lifetime of the ndarray, such that weights are still valid when functional API returned.

(self, ndarray: np.ndarray)

        layer.metadata = layer.name

    def register_ndarray(self, ndarray: np.ndarray) -> None:
        ''' When the functional APIs need to create local numpy array and use as weights for constant or other layers,
            they need to register the ndarray objects to the TRT-LLM Network to prolong the lifetime of the ndarray, such that weights are
            still valid when functional API returned.
            All the weights referenced by the trt Network are weak referenced, it's TRT-LLM's responsibility to keep the weights alive
            during the TRT network construction and TRT engine building process.
        '''
        self._registered_ndarrays.append(ndarray)

    def _generate_optimization_profiles(self) -> List[trt.IOptimizationProfile]:
        input_tensors = self._inputs