(self)
| 1302 | ) |
| 1303 | |
| 1304 | def test_multi_state_plan(self) -> None: |
| 1305 | eager_module = MultiEntryPointStatefulModel().eval() |
| 1306 | forward = export(eager_module, eager_module.get_example_inputs()) |
| 1307 | with patch_forward(eager_module, eager_module.get_state): |
| 1308 | get_state = export(eager_module, ()) |
| 1309 | with patch_forward(eager_module, eager_module.set_state): |
| 1310 | set_state = export(eager_module, (torch.zeros(1),)) |
| 1311 | edge = to_edge( |
| 1312 | {"forward": forward, "set_state": set_state, "get_state": get_state} |
| 1313 | ) |
| 1314 | et = edge.to_executorch( |
| 1315 | ExecutorchBackendConfig( |
| 1316 | memory_planning_pass=MemoryPlanningPass(share_mutable_buffers=True), |
| 1317 | emit_mutable_buffer_names=True, |
| 1318 | ) |
| 1319 | ) |
| 1320 | et_prog = et.executorch_program |
| 1321 | count = 0 |
| 1322 | for plan in et_prog.execution_plan: |
| 1323 | for value in plan.values: |
| 1324 | if ( |
| 1325 | hasattr(value.val, "allocation_info") |
| 1326 | and value.val.allocation_info is not None |
| 1327 | and value.val.allocation_info.memory_id == 2 |
| 1328 | ): |
| 1329 | count += 1 |
| 1330 | self.assertEqual(value.val.allocation_info.memory_offset_low, 0) |
| 1331 | self.assertTrue(value.val.extra_tensor_info is not None) |
| 1332 | self.assertEqual( |
| 1333 | value.val.extra_tensor_info.fully_qualified_name, "state" |
| 1334 | ) |
| 1335 | self.assertEqual(count, 3) |
| 1336 | |
| 1337 | def test_custom_kv_cache_shared_buffers(self) -> None: |
| 1338 | from executorch.examples.models.llama.source_transformation.custom_kv_cache import ( |
nothing calls this directly
no test coverage detected