| 655 | |
| 656 | args = alloc.get_unique_args() |
| 657 | assert len(args) == 5 |
| 658 | assert args[0].shape == (4,) |
| 659 | assert args[1] == 99 |
| 660 | assert args[2].shape == (2, 3) |
| 661 | assert args[3].shape == (5,) |
| 662 | assert args[4].shape == (2, 2) |
| 663 | # Outputs are zeroed |
| 664 | assert torch.equal(args[3], torch.zeros(5)) |
| 665 | assert torch.equal(args[4], torch.zeros(2, 2)) |
| 666 | |
| 667 | def test_empty_inputs_and_outputs(self): |
| 668 | """Works with no inputs and no outputs.""" |
| 669 | alloc = ShiftingMemoryPoolAllocator([], [], total_iterations=1) |
| 670 | assert alloc.get_unique_args() == [] |
| 671 | |
| 672 | def test_scalar_tensor_input(self): |
| 673 | """0-dim (scalar) tensors are handled correctly.""" |
| 674 | src = torch.tensor(5.0) |
| 675 | alloc = ShiftingMemoryPoolAllocator([src], [], total_iterations=3) |
| 676 | |
| 677 | ptrs = [] |
| 678 | for _ in range(3): |
| 679 | args = alloc.get_unique_args() |
| 680 | assert args[0].shape == () |
| 681 | assert args[0].item() == 5.0 |
| 682 | ptrs.append(args[0].data_ptr()) |
| 683 | assert len(set(ptrs)) == 3 |
| 684 | |
| 685 | |
| 686 | # ------------------------------------------------------------------ |
| 687 | # Shared helpers for gen_inputs / load_safetensors / normalize_outputs |
| 688 | # ------------------------------------------------------------------ |
| 689 | |
| 690 | _REFERENCE = "def run(a): return a" |
| 691 | |
| 692 | |
| 693 | def _make_definition(**overrides): |
| 694 | base = dict( |
| 695 | name="test_op", |
| 696 | op_type="test", |
| 697 | axes={"N": {"type": "var"}}, |
| 698 | inputs={"a": {"shape": ["N"], "dtype": "float32"}}, |
| 699 | outputs={"b": {"shape": ["N"], "dtype": "float32"}}, |
| 700 | reference=_REFERENCE, |
| 701 | ) |
| 702 | base.update(overrides) |
| 703 | return Definition(**base) |
| 704 | |
| 705 | |
| 706 | def _make_workload(**overrides): |
| 707 | base = dict(uuid="test-uuid", axes={"N": 4}, inputs={"a": {"type": "random"}}) |
| 708 | base.update(overrides) |
| 709 | return Workload(**base) |
| 710 | |
| 711 | |
| 712 | # ------------------------------------------------------------------ |
| 713 | # normalize_outputs |
| 714 | # ------------------------------------------------------------------ |
nothing calls this directly
no outgoing calls
no test coverage detected