Star857
1350
代码
代码
Issues1663
Pull Requests2223
项目讨论
Wiki
分析
Star857
1350
已合并
test(jit): add ScriptModule API alignment test cases [v2.7.1] #37630
TensorLake创建于 6月4日
test(jit): add ScriptModule API alignment test cases [v2.7.1] #37630
已合并
TensorLake创建于 6月4日
从TensorLake/torch-nputest/script-module-npu-v2.7.1合入到Ascend/pytorchv2.7.1
共 1 个文件变更+791-0
Atest/jit/test_script_module.py+791-0
@@ -0,0 +1,791 @@
1+"""
2+Add validation cases for torch.jit.ScriptModule APIs on NPU:
3+1. PyTorch community lacks sufficient direct API validations for
4+ ScriptModule instance methods, so this file is added.
5+2. This file validates 19 torch.jit.ScriptModule APIs:
6+ train, eval, requires_grad_, zero_grad,
7+ float, double, to, type, to_empty, xpu,
8+ save, state_dict, set_extra_state, share_memory,
9+ register_module, register_parameter, set_submodule,
10+ get_buffer, extra_repr (extendable).
11+"""
12+ 
13+import io
14+import os
15+import re
16+import tempfile
17+ 
18+import torch
19+import torch.nn as nn
20+from torch.testing._internal.common_utils import run_tests, TestCase
21+ 
22+ 
23+device_type = acc.type if (acc := torch.accelerator.current_accelerator()) else "cpu"
24+ 
25+ 
26+class TestScriptModuleTraining(TestCase):
27+ 
28+ def test_train_sets_training_mode(self):
29+ class M(torch.jit.ScriptModule):
30+ def __init__(self):
31+ super().__init__()
32+ self.linear = nn.Linear(2, 2)
33+ 
34+ @torch.jit.script_method
35+ def forward(self, x):
36+ return self.linear(x)
37+ 
38+ model = M().to(device_type)
39+ model.train()
40+ self.assertTrue(model.training)
41+ 
42+ def test_train_returns_self(self):
43+ class M(torch.jit.ScriptModule):
44+ def __init__(self):
45+ super().__init__()
46+ self.linear = nn.Linear(2, 2)
47+ 
48+ @torch.jit.script_method
49+ def forward(self, x):
50+ return self.linear(x)
51+ 
52+ model = M().to(device_type)
53+ result = model.train()
54+ self.assertIs(result, model)
55+ 
56+ def test_train_chained_call(self):
57+ class M(torch.jit.ScriptModule):
58+ def __init__(self):
59+ super().__init__()
60+ self.linear = nn.Linear(2, 2)
61+ 
62+ @torch.jit.script_method
63+ def forward(self, x):
64+ return self.linear(x)
65+ 
66+ model = M().to(device_type)
67+ result = model.train().train()
68+ self.assertTrue(result.training)
69+ 
70+ def test_eval_sets_eval_mode(self):
71+ class M(torch.jit.ScriptModule):
72+ def __init__(self):
73+ super().__init__()
74+ self.linear = nn.Linear(2, 2)
75+ 
76+ @torch.jit.script_method
77+ def forward(self, x):
78+ return self.linear(x)
79+ 
80+ model = M().to(device_type)
81+ model.eval()
82+ self.assertFalse(model.training)
83+ 
84+ def test_eval_returns_self(self):
85+ class M(torch.jit.ScriptModule):
86+ def __init__(self):
87+ super().__init__()
88+ self.linear = nn.Linear(2, 2)
89+ 
90+ @torch.jit.script_method
91+ def forward(self, x):
92+ return self.linear(x)
93+ 
94+ model = M().to(device_type)
95+ result = model.eval()
96+ self.assertIs(result, model)
97+ 
98+ def test_eval_chained_call(self):
99+ class M(torch.jit.ScriptModule):
100+ def __init__(self):
101+ super().__init__()
102+ self.linear = nn.Linear(2, 2)
103+ 
104+ @torch.jit.script_method
105+ def forward(self, x):
106+ return self.linear(x)
107+ 
108+ model = M().to(device_type)
109+ result = model.eval().eval()
110+ self.assertFalse(result.training)
111+ 
112+ def test_train_eval_toggle(self):
113+ class M(torch.jit.ScriptModule):
114+ def __init__(self):
115+ super().__init__()
116+ self.linear = nn.Linear(2, 2)
117+ 
118+ @torch.jit.script_method
119+ def forward(self, x):
120+ return self.linear(x)
121+ 
122+ model = M().to(device_type)
123+ model.train()
124+ self.assertTrue(model.training)
125+ model.eval()
126+ self.assertFalse(model.training)
127+ model.train()
128+ self.assertTrue(model.training)
129+ 
130+ def test_train_propagates_to_submodule(self):
131+ class Sub(torch.jit.ScriptModule):
132+ def __init__(self):
133+ super().__init__()
134+ self.linear = nn.Linear(2, 2)
135+ 
136+ @torch.jit.script_method
137+ def forward(self, x):
138+ return self.linear(x)
139+ 
140+ class M(torch.jit.ScriptModule):
141+ def __init__(self):
142+ super().__init__()
143+ self.sub = Sub()
144+ 
145+ @torch.jit.script_method
146+ def forward(self, x):
147+ return self.sub(x)
148+ 
149+ model = M().to(device_type)
150+ model.eval()
151+ self.assertFalse(model.sub.training)
152+ model.train()
153+ self.assertTrue(model.sub.training)
154+ 
155+ def test_requires_grad_sets_flag(self):
156+ class M(torch.jit.ScriptModule):
157+ def __init__(self):
158+ super().__init__()
159+ self.linear = nn.Linear(2, 2)
160+ 
161+ @torch.jit.script_method
162+ def forward(self, x):
163+ return self.linear(x)
164+ 
165+ model = M().to(device_type)
166+ model.requires_grad_(True)
167+ self.assertTrue(model.linear.weight.requires_grad)
168+ 
169+ def test_requires_grad_returns_self(self):
170+ class M(torch.jit.ScriptModule):
171+ def __init__(self):
172+ super().__init__()
173+ self.linear = nn.Linear(2, 2)
174+ 
175+ @torch.jit.script_method
176+ def forward(self, x):
177+ return self.linear(x)
178+ 
179+ model = M().to(device_type)
180+ result = model.requires_grad_(False)
181+ self.assertIs(result, model)
182+ 
183+ def test_requires_grad_false(self):
184+ class M(torch.jit.ScriptModule):
185+ def __init__(self):
186+ super().__init__()
187+ self.linear = nn.Linear(2, 2)
188+ 
189+ @torch.jit.script_method
190+ def forward(self, x):
191+ return self.linear(x)
192+ 
193+ model = M().to(device_type)
194+ model.requires_grad_(False)
195+ self.assertFalse(model.linear.weight.requires_grad)
196+ 
197+ def test_zero_grad_clears_gradients(self):
198+ class M(torch.jit.ScriptModule):
199+ def __init__(self):
200+ super().__init__()
201+ self.linear = nn.Linear(2, 2)
202+ 
203+ @torch.jit.script_method
204+ def forward(self, x):
205+ return self.linear(x)
206+ 
207+ model = M().to(device_type)
208+ x = torch.randn(2, 2, requires_grad=True).to(device_type)
209+ out = model(x)
210+ out.sum().backward()
211+ self.assertIsNotNone(model.linear.weight.grad)
212+ model.zero_grad()
213+ self.assertIsNone(model.linear.weight.grad)
214+ 
215+ def test_zero_grad_set_to_none_false(self):
216+ class M(torch.jit.ScriptModule):
217+ def __init__(self):
218+ super().__init__()
219+ self.linear = nn.Linear(2, 2)
220+ 
221+ @torch.jit.script_method
222+ def forward(self, x):
223+ return self.linear(x)
224+ 
225+ model = M().to(device_type)
226+ x = torch.randn(2, 2, requires_grad=True).to(device_type)
227+ out = model(x)
228+ out.sum().backward()
229+ self.assertIsNotNone(model.linear.weight.grad)
230+ model.zero_grad(set_to_none=False)
231+ self.assertIsNotNone(model.linear.weight.grad)
232+ self.assertEqual(model.linear.weight.grad, torch.zeros_like(
233+ model.linear.weight.grad))
234+ 
235+ def test_zero_grad_no_gradients(self):
236+ class M(torch.jit.ScriptModule):
237+ def __init__(self):
238+ super().__init__()
239+ self.linear = nn.Linear(2, 2)
240+ 
241+ @torch.jit.script_method
242+ def forward(self, x):
243+ return self.linear(x)
244+ 
245+ model = M().to(device_type)
246+ model.zero_grad()
247+ 
248+ 
249+class TestScriptModuleDtype(TestCase):
250+ 
251+ def test_float_converts_parameters(self):
252+ class M(torch.jit.ScriptModule):
253+ def __init__(self):
254+ super().__init__()
255+ self.linear = nn.Linear(2, 2)
256+ 
257+ @torch.jit.script_method
258+ def forward(self, x):
259+ return self.linear(x)
260+ 
261+ model = M().to(device_type)
262+ model.float()
263+ self.assertEqual(model.linear.weight.dtype, torch.float32)
264+ 
265+ def test_float_returns_self(self):
266+ class M(torch.jit.ScriptModule):
267+ def __init__(self):
268+ super().__init__()
269+ self.linear = nn.Linear(2, 2)
270+ 
271+ @torch.jit.script_method
272+ def forward(self, x):
273+ return self.linear(x)
274+ 
275+ model = M().to(device_type)
276+ result = model.float()
277+ self.assertIs(result, model)
278+ 
279+ def test_double_converts_parameters(self):
280+ class M(torch.jit.ScriptModule):
281+ def __init__(self):
282+ super().__init__()
283+ self.linear = nn.Linear(2, 2)
284+ 
285+ @torch.jit.script_method
286+ def forward(self, x):
287+ return self.linear(x)
288+ 
289+ model = M().to(device_type)
290+ model.double()
291+ # NPU may cast double to float, verify actual dtype
292+ self.assertIn(model.linear.weight.dtype, (torch.float64, torch.float32))
293+ 
294+ def test_double_returns_self(self):
295+ class M(torch.jit.ScriptModule):
296+ def __init__(self):
297+ super().__init__()
298+ self.linear = nn.Linear(2, 2)
299+ 
300+ @torch.jit.script_method
301+ def forward(self, x):
302+ return self.linear(x)
303+ 
304+ model = M().to(device_type)
305+ result = model.double()
306+ self.assertIs(result, model)
307+ 
308+ def test_to_dtype(self):
309+ class M(torch.jit.ScriptModule):
310+ def __init__(self):
311+ super().__init__()
312+ self.linear = nn.Linear(2, 2)
313+ 
314+ @torch.jit.script_method
315+ def forward(self, x):
316+ return self.linear(x)
317+ 
318+ model = M().to(device_type)
319+ model.to(torch.float64)
320+ # NPU may cast double to float
321+ self.assertIn(model.linear.weight.dtype, (torch.float64, torch.float32))
322+ 
323+ def test_to_device(self):
324+ class M(torch.jit.ScriptModule):
325+ def __init__(self):
326+ super().__init__()
327+ self.linear = nn.Linear(2, 2)
328+ 
329+ @torch.jit.script_method
330+ def forward(self, x):
331+ return self.linear(x)
332+ 
333+ model = M().to(device_type)
334+ model.to(device_type)
335+ self.assertEqual(model.linear.weight.device.type, device_type)
336+ 
337+ def test_to_returns_self_or_copy(self):
338+ class M(torch.jit.ScriptModule):
339+ def __init__(self):
340+ super().__init__()
341+ self.linear = nn.Linear(2, 2)
342+ 
343+ @torch.jit.script_method
344+ def forward(self, x):
345+ return self.linear(x)
346+ 
347+ model = M().to(device_type)
348+ result = model.to(device_type)
349+ self.assertIsInstance(result, torch.jit.ScriptModule)
350+ 
351+ def test_to_chained_call(self):
352+ class M(torch.jit.ScriptModule):
353+ def __init__(self):
354+ super().__init__()
355+ self.linear = nn.Linear(2, 2)
356+ 
357+ @torch.jit.script_method
358+ def forward(self, x):
359+ return self.linear(x)
360+ 
361+ model = M().to(device_type)
362+ result = model.to(torch.float64).to(torch.float32)
363+ self.assertEqual(model.linear.weight.dtype, torch.float32)
364+ 
365+ def test_type_float32_no_downgrade(self):
366+ class M(torch.jit.ScriptModule):
367+ def __init__(self):
368+ super().__init__()
369+ self.linear = nn.Linear(2, 2)
370+ 
371+ @torch.jit.script_method
372+ def forward(self, x):
373+ return self.linear(x)
374+ 
375+ model = M().to(device_type)
376+ # Normal path: float32 input stays float32
377+ model.type(torch.float32)
378+ self.assertEqual(model.linear.weight.dtype, torch.float32)
379+ 
380+ def test_type_float64(self):
381+ class M(torch.jit.ScriptModule):
382+ def __init__(self):
383+ super().__init__()
384+ self.linear = nn.Linear(2, 2)
385+ 
386+ @torch.jit.script_method
387+ def forward(self, x):
388+ return self.linear(x)
389+ 
390+ model = M().to(device_type)
391+ model.type(torch.float64)
392+ # NPU may cast double to float
393+ self.assertIn(model.linear.weight.dtype, (torch.float64, torch.float32))
394+ 
395+ def test_type_returns_self(self):
396+ class M(torch.jit.ScriptModule):
397+ def __init__(self):
398+ super().__init__()
399+ self.linear = nn.Linear(2, 2)
400+ 
401+ @torch.jit.script_method
402+ def forward(self, x):
403+ return self.linear(x)
404+ 
405+ model = M().to(device_type)
406+ result = model.type(torch.float32)
407+ self.assertIs(result, model)
408+ 
409+ def test_to_empty_moves_to_device(self):
410+ class M(torch.jit.ScriptModule):
411+ def __init__(self):
412+ super().__init__()
413+ self.linear = nn.Linear(2, 2)
414+ 
415+ @torch.jit.script_method
416+ def forward(self, x):
417+ return self.linear(x)
418+ 
419+ model = M().to(device_type)
420+ result = model.to_empty(device=device_type)
421+ self.assertIsInstance(result, torch.jit.ScriptModule)
422+ self.assertEqual(result.linear.weight.device.type, device_type)
423+ 
424+ def test_to_empty_returns_self(self):
425+ class M(torch.jit.ScriptModule):
426+ def __init__(self):
427+ super().__init__()
428+ self.linear = nn.Linear(2, 2)
429+ 
430+ @torch.jit.script_method
431+ def forward(self, x):
432+ return self.linear(x)
433+ 
434+ model = M().to(device_type)
435+ result = model.to_empty(device=device_type)
436+ self.assertIs(result, model)
437+ 
438+ def test_xpu_raises(self):
439+ class M(torch.jit.ScriptModule):
440+ def __init__(self):
441+ super().__init__()
442+ self.linear = nn.Linear(2, 2)
443+ 
444+ @torch.jit.script_method
445+ def forward(self, x):
446+ return self.linear(x)
447+ 
448+ model = M().to(device_type)
449+ # XPU is not compiled in current environment
450+ with self.assertRaises(AssertionError):
451+ model.xpu()
452+ 
453+ 
454+class TestScriptModuleSerialization(TestCase):
455+ 
456+ def test_save_to_file(self):
457+ class M(torch.jit.ScriptModule):
458+ def __init__(self):
459+ super().__init__()
460+ self.linear = nn.Linear(2, 2)
461+ 
462+ @torch.jit.script_method
463+ def forward(self, x):
464+ return self.linear(x)
465+ 
466+ model = M().to(device_type)
467+ with tempfile.NamedTemporaryFile(suffix=".pt", delete=False) as f:
468+ path = f.name
469+ try:
470+ model.save(path)
471+ self.assertTrue(os.path.exists(path))
472+ self.assertGreater(os.path.getsize(path), 0)
473+ finally:
474+ if os.path.exists(path):
475+ os.remove(path)
476+ 
477+ def test_save_and_load(self):
478+ class M(torch.jit.ScriptModule):
479+ def __init__(self):
480+ super().__init__()
481+ self.linear = nn.Linear(2, 2)
482+ 
483+ @torch.jit.script_method
484+ def forward(self, x):
485+ return self.linear(x)
486+ 
487+ model = M().to(device_type)
488+ with tempfile.NamedTemporaryFile(suffix=".pt", delete=False) as f:
489+ path = f.name
490+ try:
491+ model.save(path)
492+ loaded = torch.jit.load(path)
493+ x = torch.randn(2, 2).to(device_type)
494+ self.assertEqual(model(x), loaded(x))
495+ finally:
496+ if os.path.exists(path):
497+ os.remove(path)
498+ 
499+ def test_save_preserves_parameters(self):
500+ class M(torch.jit.ScriptModule):
501+ def __init__(self):
502+ super().__init__()
503+ self.linear = nn.Linear(2, 2)
504+ 
505+ @torch.jit.script_method
506+ def forward(self, x):
507+ return self.linear(x)
508+ 
509+ model = M().to(device_type)
510+ with torch.no_grad():
511+ model.linear.weight.fill_(1.0)
512+ model.linear.bias.fill_(2.0)
513+ 
514+ with tempfile.NamedTemporaryFile(suffix=".pt", delete=False) as f:
515+ path = f.name
516+ try:
517+ model.save(path)
518+ loaded = torch.jit.load(path)
519+ self.assertEqual(loaded.linear.weight, torch.ones(2, 2))
520+ self.assertEqual(loaded.linear.bias, torch.ones(2) * 2)
521+ finally:
522+ if os.path.exists(path):
523+ os.remove(path)
524+ 
525+ def test_save_to_buffer(self):
526+ class M(torch.jit.ScriptModule):
527+ def __init__(self):
528+ super().__init__()
529+ self.linear = nn.Linear(2, 2)
530+ 
531+ @torch.jit.script_method
532+ def forward(self, x):
533+ return self.linear(x)
534+ 
535+ model = M().to(device_type)
536+ buffer = model.save_to_buffer()
537+ self.assertIsInstance(buffer, bytes)
538+ self.assertGreater(len(buffer), 0)
539+ 
540+ def test_save_to_buffer_and_load(self):
541+ class M(torch.jit.ScriptModule):
542+ def __init__(self):
543+ super().__init__()
544+ self.linear = nn.Linear(2, 2)
545+ 
546+ @torch.jit.script_method
547+ def forward(self, x):
548+ return self.linear(x)
549+ 
550+ model = M().to(device_type)
551+ buffer = model.save_to_buffer()
552+ loaded = torch.jit.load(io.BytesIO(buffer))
553+ x = torch.randn(2, 2).to(device_type)
554+ self.assertEqual(model(x), loaded(x))
555+ 
556+ def test_state_dict_returns_dict(self):
557+ class M(torch.jit.ScriptModule):
558+ def __init__(self):
559+ super().__init__()
560+ self.linear = nn.Linear(2, 2)
561+ 
562+ @torch.jit.script_method
563+ def forward(self, x):
564+ return self.linear(x)
565+ 
566+ model = M().to(device_type)
567+ sd = model.state_dict()
568+ self.assertIsInstance(sd, dict)
569+ self.assertIn("linear.weight", sd)
570+ self.assertIn("linear.bias", sd)
571+ 
572+ def test_state_dict_values_on_npu(self):
573+ class M(torch.jit.ScriptModule):
574+ def __init__(self):
575+ super().__init__()
576+ self.linear = nn.Linear(2, 2)
577+ 
578+ @torch.jit.script_method
579+ def forward(self, x):
580+ return self.linear(x)
581+ 
582+ model = M().to(device_type)
583+ sd = model.state_dict()
584+ self.assertEqual(sd["linear.weight"].device.type, device_type)
585+ 
586+ def test_state_dict_no_training_key(self):
587+ class M(torch.jit.ScriptModule):
588+ def __init__(self):
589+ super().__init__()
590+ self.linear = nn.Linear(2, 2)
591+ 
592+ @torch.jit.script_method
593+ def forward(self, x):
594+ return self.linear(x)
595+ 
596+ model = M().to(device_type)
597+ sd = model.state_dict()
598+ self.assertNotIn("training", sd)
599+ 
600+ def test_state_dict_buffer_included(self):
601+ class M(torch.jit.ScriptModule):
602+ def __init__(self):
603+ super().__init__()
604+ self.register_buffer("buf", torch.ones(2, 2))
605+ 
606+ @torch.jit.script_method
607+ def forward(self, x):
608+ return x + self.buf
609+ 
610+ model = M().to(device_type)
611+ sd = model.state_dict()
612+ self.assertIn("buf", sd)
613+ self.assertEqual(sd["buf"], torch.ones(2, 2))
614+ 
615+ def test_set_extra_state_raises(self):
616+ class M(torch.jit.ScriptModule):
617+ def __init__(self):
618+ super().__init__()
619+ self.linear = nn.Linear(2, 2)
620+ 
621+ @torch.jit.script_method
622+ def forward(self, x):
623+ return self.linear(x)
624+ 
625+ model = M().to(device_type)
626+ with self.assertRaises(RuntimeError):
627+ model.set_extra_state({"version": 1})
628+ 
629+ def test_share_memory_raises_on_npu(self):
630+ class M(torch.jit.ScriptModule):
631+ def __init__(self):
632+ super().__init__()
633+ self.linear = nn.Linear(2, 2)
634+ 
635+ @torch.jit.script_method
636+ def forward(self, x):
637+ return self.linear(x)
638+ 
639+ model = M().to(device_type)
640+ # share_memory is intercepted by torch-npu for NPU modules
641+ with self.assertRaises(RuntimeError):
642+ model.share_memory()
643+ 
644+ 
645+class TestScriptModuleMetadata(TestCase):
646+ 
647+ def test_register_module_raises_on_npu(self):
648+ class M(torch.jit.ScriptModule):
649+ def __init__(self):
650+ super().__init__()
651+ self.linear = nn.Linear(2, 2)
652+ 
653+ @torch.jit.script_method
654+ def forward(self, x):
655+ return self.linear(x)
656+ 
657+ model = M().to(device_type)
658+ sub = nn.Linear(2, 2).to(device_type)
659+ # torch-npu intercepts register_module for NPU modules
660+ with self.assertRaises(RuntimeError):
661+ model.register_module("new_sub", sub)
662+ 
663+ def test_register_parameter_raises_on_npu(self):
664+ class M(torch.jit.ScriptModule):
665+ def __init__(self):
666+ super().__init__()
667+ self.linear = nn.Linear(2, 2)
668+ 
669+ @torch.jit.script_method
670+ def forward(self, x):
671+ return self.linear(x)
672+ 
673+ model = M().to(device_type)
674+ param = nn.Parameter(torch.randn(2, 2)).to(device_type)
675+ # torch-npu intercepts register_parameter for NPU modules
676+ with self.assertRaises(RuntimeError):
677+ model.register_parameter("new_param", param)
678+ 
679+ def test_set_submodule_raises(self):
680+ class M(torch.jit.ScriptModule):
681+ def __init__(self):
682+ super().__init__()
683+ self.linear = nn.Linear(2, 2)
684+ 
685+ @torch.jit.script_method
686+ def forward(self, x):
687+ return self.linear(x)
688+ 
689+ model = M().to(device_type)
690+ sub = nn.Linear(2, 2).to(device_type)
691+ # Cannot re-assign modules in a ScriptModule
692+ with self.assertRaises(RuntimeError):
693+ model.set_submodule("linear", sub)
694+ 
695+ def test_set_submodule_nested_raises(self):
696+ class Sub(torch.jit.ScriptModule):
697+ def __init__(self):
698+ super().__init__()
699+ self.linear = nn.Linear(2, 2)
700+ 
701+ @torch.jit.script_method
702+ def forward(self, x):
703+ return self.linear(x)
704+ 
705+ class M(torch.jit.ScriptModule):
706+ def __init__(self):
707+ super().__init__()
708+ self.sub = Sub()
709+ 
710+ @torch.jit.script_method
711+ def forward(self, x):
712+ return self.sub(x)
713+ 
714+ model = M().to(device_type)
715+ new_sub = nn.Linear(2, 2).to(device_type)
716+ with self.assertRaises(RuntimeError):
717+ model.set_submodule("sub.linear", new_sub)
718+ 
719+ def test_get_buffer_returns_buffer(self):
720+ class M(torch.jit.ScriptModule):
721+ def __init__(self):
722+ super().__init__()
723+ self.register_buffer("buf", torch.ones(2, 2))
724+ 
725+ @torch.jit.script_method
726+ def forward(self, x):
727+ return x + self.buf
728+ 
729+ model = M().to(device_type)
730+ buf = model.get_buffer("buf")
731+ self.assertEqual(buf, torch.ones(2, 2))
732+ 
733+ def test_get_buffer_nested(self):
734+ class Sub(torch.jit.ScriptModule):
735+ def __init__(self):
736+ super().__init__()
737+ self.register_buffer("buf", torch.ones(2, 2))
738+ 
739+ @torch.jit.script_method
740+ def forward(self, x):
741+ return x + self.buf
742+ 
743+ class M(torch.jit.ScriptModule):
744+ def __init__(self):
745+ super().__init__()
746+ self.sub = Sub()
747+ 
748+ @torch.jit.script_method
749+ def forward(self, x):
750+ return self.sub(x)
751+ 
752+ model = M().to(device_type)
753+ buf = model.get_buffer("sub.buf")
754+ self.assertEqual(buf, torch.ones(2, 2))
755+ 
756+ def test_extra_repr_returns_string(self):
757+ class MyScriptModule(torch.jit.ScriptModule):
758+ def __init__(self):
759+ super().__init__()
760+ self.linear = nn.Linear(2, 2)
761+ 
762+ @torch.jit.script_method
763+ def forward(self, x):
764+ return self.linear(x)
765+ 
766+ model = MyScriptModule().to(device_type)
767+ result = model.extra_repr()
768+ self.assertIsInstance(result, str)
769+ # PyTorch 2.7.1 returns empty string on ScriptModule extra_repr.
770+ # Verify that str(model) contains ScriptModule class info.
771+ full_repr = str(model)
772+ self.assertIn("ScriptModule", full_repr)
773+ 
774+ def test_extra_repr_matches_pattern(self):
775+ class MyScriptModule(torch.jit.ScriptModule):
776+ def __init__(self):
777+ super().__init__()
778+ self.linear = nn.Linear(2, 2)
779+ 
780+ @torch.jit.script_method
781+ def forward(self, x):
782+ return self.linear(x)
783+ 
784+ model = MyScriptModule().to(device_type)
785+ # Verify the model string contains a class-like identifier
786+ full_repr = repr(model)
787+ self.assertIsNotNone(re.search(r"MyScriptModule|ScriptModule", full_repr))
788+ 
789+ 
790+if __name__ == "__main__":
791+ run_tests()