Home
last modified time | relevance | path

Searched refs:engine_input_shapes (Results 1 – 2 of 2) sorted by relevance

/aosp_15_r20/external/tensorflow/tensorflow/compiler/tf2tensorrt/kernels/
H A Dtrt_engine_resource_ops.cc129 std::vector<TensorShape> engine_input_shapes; in Compute() local
131 engine_input_shapes.reserve(input_shapes.size()); in Compute()
133 engine_input_shapes.emplace_back(shape); in Compute()
150 raw_engine, engine_input_shapes.size())); in Compute()
158 resource->cache_.emplace(engine_input_shapes, in Compute()
232 const std::vector<TensorShape>& engine_input_shapes = pair.first; in Compute() local
233 for (const TensorShape& shape : engine_input_shapes) { in Compute()
258 for (const TensorShape& shape : engine_input_shapes) { in Compute()
H A Dtrt_engine_op.cc1184 std::vector<TensorShape> engine_input_shapes(input_concrete_shapes); in GetEngine() local
1191 for (int i = 0; i < engine_input_shapes.size(); i++) { in GetEngine()
1192 engine_input_shapes[i].set_dim(0, max_batch_size); in GetEngine()
1199 cache.emplace(engine_input_shapes, in GetEngine()
1211 return std::pair<EngineContext*, int>(cache.at(engine_input_shapes).get(), in GetEngine()