|
27 | 27 | Program, |
28 | 28 | ) |
29 | 29 | from executorch.exir import to_edge_transform_and_lower |
30 | | -from executorch.exir._serialize._ptn import ALIASES_ENTRY, PTG_ENTRY, SAFETENSORS_ENTRY |
| 30 | +from executorch.exir._serialize._ptn import ( |
| 31 | + ALIASES_ENTRY, |
| 32 | + PTG_ENTRY, |
| 33 | + read_ptn, |
| 34 | + SAFETENSORS_ENTRY, |
| 35 | +) |
31 | 36 | from executorch.exir.lowered_backend_module import ( |
32 | 37 | executorch_call_delegate, |
33 | 38 | get_lowered_submodules, |
@@ -393,6 +398,58 @@ def test_dropped_public_input_is_rejected(self): |
393 | 398 | self._check(self._graph_module(drop_input=True)) |
394 | 399 |
|
395 | 400 |
|
| 401 | +class LiftedConstantCollisionTest(unittest.TestCase): |
| 402 | + def _scalar_shift(self, offset: float): |
| 403 | + class Shift(nn.Module): |
| 404 | + def forward(self, x): |
| 405 | + return x + offset |
| 406 | + |
| 407 | + return torch.export.export(Shift(), (torch.randn(2, 4),)) |
| 408 | + |
| 409 | + def _read(self, native) -> tuple[Program, dict[str, torch.Tensor]]: |
| 410 | + with tempfile.TemporaryDirectory() as td: |
| 411 | + path = os.path.join(td, "m.ptn") |
| 412 | + native.save(path) |
| 413 | + ptg, constants = read_ptn(path) |
| 414 | + return deserialize_program(ptg), constants |
| 415 | + |
| 416 | + def _lifted_keys(self, program: Program, method_name: str) -> list[str]: |
| 417 | + (method,) = [m for m in program.methods if m.name == method_name] |
| 418 | + return [c.data_key for c in method.constants or [] if "lifted" in c.data_key] |
| 419 | + |
| 420 | + def test_distinct_lifted_constants_are_namespaced(self): |
| 421 | + programs = {"a": self._scalar_shift(1.0), "b": self._scalar_shift(2.0)} |
| 422 | + native = to_native(programs) |
| 423 | + self.assertEqual(native.methods, {"a", "b"}) |
| 424 | + program, constants = self._read(native) |
| 425 | + (a_key,) = self._lifted_keys(program, "a") |
| 426 | + (b_key,) = self._lifted_keys(program, "b") |
| 427 | + self.assertEqual(b_key, f"b/{a_key}") |
| 428 | + torch.testing.assert_close( |
| 429 | + constants[a_key], torch.full_like(constants[a_key], 1.0), rtol=0, atol=0 |
| 430 | + ) |
| 431 | + torch.testing.assert_close( |
| 432 | + constants[b_key], torch.full_like(constants[b_key], 2.0), rtol=0, atol=0 |
| 433 | + ) |
| 434 | + |
| 435 | + def test_identical_lifted_constants_stay_shared(self): |
| 436 | + programs = {"a": self._scalar_shift(1.0), "b": self._scalar_shift(1.0)} |
| 437 | + program, constants = self._read(to_native(programs)) |
| 438 | + self.assertEqual( |
| 439 | + self._lifted_keys(program, "a"), self._lifted_keys(program, "b") |
| 440 | + ) |
| 441 | + lifted = {k for k in constants if "lifted" in k} |
| 442 | + self.assertEqual(len(lifted), 1, lifted) |
| 443 | + |
| 444 | + def test_parameter_fqn_collision_is_still_rejected(self): |
| 445 | + programs = { |
| 446 | + "a": torch.export.export(nn.Linear(4, 4), (torch.randn(1, 4),)), |
| 447 | + "b": torch.export.export(nn.Linear(4, 4), (torch.randn(1, 4),)), |
| 448 | + } |
| 449 | + with self.assertRaisesRegex(ValueError, "differs between methods"): |
| 450 | + to_native(programs) |
| 451 | + |
| 452 | + |
396 | 453 | class MutableBufferDedupTest(unittest.TestCase): |
397 | 454 | def test_identical_mutable_buffers_stay_separate(self): |
398 | 455 | """Equal bytes must not make two independently mutable buffers one. |
|
0 commit comments