forked from pytorch/pytorch.github.io
-
Notifications
You must be signed in to change notification settings - Fork 0
/
Copy pathjit_language_reference_v2.html
2692 lines (2431 loc) · 243 KB
/
jit_language_reference_v2.html
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
704
705
706
707
708
709
710
711
712
713
714
715
716
717
718
719
720
721
722
723
724
725
726
727
728
729
730
731
732
733
734
735
736
737
738
739
740
741
742
743
744
745
746
747
748
749
750
751
752
753
754
755
756
757
758
759
760
761
762
763
764
765
766
767
768
769
770
771
772
773
774
775
776
777
778
779
780
781
782
783
784
785
786
787
788
789
790
791
792
793
794
795
796
797
798
799
800
801
802
803
804
805
806
807
808
809
810
811
812
813
814
815
816
817
818
819
820
821
822
823
824
825
826
827
828
829
830
831
832
833
834
835
836
837
838
839
840
841
842
843
844
845
846
847
848
849
850
851
852
853
854
855
856
857
858
859
860
861
862
863
864
865
866
867
868
869
870
871
872
873
874
875
876
877
878
879
880
881
882
883
884
885
886
887
888
889
890
891
892
893
894
895
896
897
898
899
900
901
902
903
904
905
906
907
908
909
910
911
912
913
914
915
916
917
918
919
920
921
922
923
924
925
926
927
928
929
930
931
932
933
934
935
936
937
938
939
940
941
942
943
944
945
946
947
948
949
950
951
952
953
954
955
956
957
958
959
960
961
962
963
964
965
966
967
968
969
970
971
972
973
974
975
976
977
978
979
980
981
982
983
984
985
986
987
988
989
990
991
992
993
994
995
996
997
998
999
1000
<!DOCTYPE html>
<!--[if IE 8]><html class="no-js lt-ie9" lang="en" > <![endif]-->
<!--[if gt IE 8]><!--> <html class="no-js" lang="en" > <!--<![endif]-->
<head>
<meta name="robots" content="noindex">
<meta charset="utf-8">
<meta name="viewport" content="width=device-width, initial-scale=1.0">
<title>TorchScript Language Reference — PyTorch 1.11.0 documentation</title>
<link rel="canonical" href="https://pytorch.org/docs/stable/jit_language_reference_v2.html"/>
<link rel="stylesheet" href="_static/css/theme.css" type="text/css" />
<!-- <link rel="stylesheet" href="_static/pygments.css" type="text/css" /> -->
<link rel="stylesheet" href="_static/copybutton.css" type="text/css" />
<link rel="stylesheet" href="https://cdn.jsdelivr.net/npm/[email protected]/dist/katex.min.css" type="text/css" />
<link rel="stylesheet" href="https://cdn.jsdelivr.net/npm/[email protected]/dist/katex.min.css" type="text/css" />
<link rel="stylesheet" href="_static/katex-math.css" type="text/css" />
<link rel="stylesheet" href="_static/css/jit.css" type="text/css" />
<link rel="index" title="Index" href="genindex.html" />
<link rel="search" title="Search" href="search.html" />
<link rel="next" title="torch.jit.script" href="generated/torch.jit.script.html" />
<link rel="prev" title="TorchScript Language Reference" href="jit_language_reference.html" />
<!-- Google Analytics -->
<script async src="https://www.googletagmanager.com/gtag/js?id=UA-117752657-2"></script>
<script>
window.dataLayer = window.dataLayer || [];
function gtag(){dataLayer.push(arguments);}
gtag('js', new Date());
gtag('config', 'UA-117752657-2');
</script>
<!-- End Google Analytics -->
<script src="_static/js/modernizr.min.js"></script>
<!-- Preload the theme fonts -->
<link rel="preload" href="_static/fonts/FreightSans/freight-sans-book.woff2" as="font" type="font/woff2" crossorigin="anonymous">
<link rel="preload" href="_static/fonts/FreightSans/freight-sans-medium.woff2" as="font" type="font/woff2" crossorigin="anonymous">
<link rel="preload" href="_static/fonts/IBMPlexMono/IBMPlexMono-Medium.woff2" as="font" type="font/woff2" crossorigin="anonymous">
<link rel="preload" href="_static/fonts/FreightSans/freight-sans-bold.woff2" as="font" type="font/woff2" crossorigin="anonymous">
<link rel="preload" href="_static/fonts/FreightSans/freight-sans-medium-italic.woff2" as="font" type="font/woff2" crossorigin="anonymous">
<link rel="preload" href="_static/fonts/IBMPlexMono/IBMPlexMono-SemiBold.woff2" as="font" type="font/woff2" crossorigin="anonymous">
<!-- Preload the katex fonts -->
<link rel="preload" href="https://cdn.jsdelivr.net/npm/[email protected]/dist/fonts/KaTeX_Math-Italic.woff2" as="font" type="font/woff2" crossorigin="anonymous">
<link rel="preload" href="https://cdn.jsdelivr.net/npm/[email protected]/dist/fonts/KaTeX_Main-Regular.woff2" as="font" type="font/woff2" crossorigin="anonymous">
<link rel="preload" href="https://cdn.jsdelivr.net/npm/[email protected]/dist/fonts/KaTeX_Main-Bold.woff2" as="font" type="font/woff2" crossorigin="anonymous">
<link rel="preload" href="https://cdn.jsdelivr.net/npm/[email protected]/dist/fonts/KaTeX_Size1-Regular.woff2" as="font" type="font/woff2" crossorigin="anonymous">
<link rel="preload" href="https://cdn.jsdelivr.net/npm/[email protected]/dist/fonts/KaTeX_Size4-Regular.woff2" as="font" type="font/woff2" crossorigin="anonymous">
<link rel="preload" href="https://cdn.jsdelivr.net/npm/[email protected]/dist/fonts/KaTeX_Size2-Regular.woff2" as="font" type="font/woff2" crossorigin="anonymous">
<link rel="preload" href="https://cdn.jsdelivr.net/npm/[email protected]/dist/fonts/KaTeX_Size3-Regular.woff2" as="font" type="font/woff2" crossorigin="anonymous">
<link rel="preload" href="https://cdn.jsdelivr.net/npm/[email protected]/dist/fonts/KaTeX_Caligraphic-Regular.woff2" as="font" type="font/woff2" crossorigin="anonymous">
<link rel="stylesheet" href="https://use.fontawesome.com/releases/v5.15.2/css/all.css" integrity="sha384-vSIIfh2YWi9wW0r9iZe7RJPrKwp6bG+s9QZMoITbCckVJqGCCRhc+ccxNcdpHuYu" crossorigin="anonymous">
</head>
<div class="container-fluid header-holder tutorials-header" id="header-holder">
<div class="container">
<div class="header-container">
<a class="header-logo" href="https://pytorch.org/" aria-label="PyTorch"></a>
<div class="main-menu">
<ul>
<li>
<a href="https://pytorch.org/get-started">Get Started</a>
</li>
<li>
<a href="https://pytorch.org/ecosystem">Ecosystem</a>
</li>
<li>
<a href="https://pytorch.org/mobile">Mobile</a>
</li>
<li>
<a href="https://pytorch.org/blog/">Blog</a>
</li>
<li>
<a href="https://pytorch.org/tutorials">Tutorials</a>
</li>
<li class="active docs-active">
<div id="resourcesDropdownButton" data-toggle="resources-dropdown" class="resources-dropdown">
<a class="resource-option with-down-orange-arrow">
Docs
</a>
<div class="resources-dropdown-menu">
<a class="doc-dropdown-option nav-dropdown-item" href="https://pytorch.org/docs/stable/index.html">
<span class="dropdown-title">PyTorch</span>
<p></p>
</a>
<a class="doc-dropdown-option nav-dropdown-item" href="https://pytorch.org/audio/stable/index.html">
<span class="dropdown-title">torchaudio</span>
<p></p>
</a>
<a class="doc-dropdown-option nav-dropdown-item" href="https://pytorch.org/text/stable/index.html">
<span class="dropdown-title">torchtext</span>
<p></p>
</a>
<a class="doc-dropdown-option nav-dropdown-item" href="https://pytorch.org/vision/stable/index.html">
<span class="dropdown-title">torchvision</span>
<p></p>
</a>
<a class="doc-dropdown-option nav-dropdown-item" href="https://pytorch.org/torchrec">
<span class="dropdown-title">TorchRec</span>
<p></p>
</a>
<a class="doc-dropdown-option nav-dropdown-item" href="https://pytorch.org/data">
<span class="dropdown-title">TorchData</span>
<p></p>
</a>
<a class="doc-dropdown-option nav-dropdown-item" href="https://pytorch.org/serve/">
<span class="dropdown-title">TorchServe</span>
<p></p>
</a>
<a class="doc-dropdown-option nav-dropdown-item" href="https://pytorch.org/xla">
<span class="dropdown-title">PyTorch on XLA Devices</span>
<p></p>
</a>
</div>
</li>
<li>
<div id="resourcesDropdownButton" data-toggle="resources-dropdown" class="resources-dropdown">
<a class="resource-option with-down-arrow">
Resources
</a>
<div class="resources-dropdown-menu">
<a class="nav-dropdown-item" href="https://pytorch.org/features">
<span class="dropdown-title">About</span>
<p>Learn about PyTorch’s features and capabilities</p>
</a>
<a class="nav-dropdown-item" href="https://pytorch.org/#community-module">
<span class="dropdown-title">Community</span>
<p>Join the PyTorch developer community to contribute, learn, and get your questions answered.</p>
</a>
<a class="nav-dropdown-item" href="https://pytorch.org/resources">
<span class="dropdown-title">Developer Resources</span>
<p>Find resources and get questions answered</p>
</a>
<a class="nav-dropdown-item" href="https://discuss.pytorch.org/" target="_blank">
<span class="dropdown-title">Forums</span>
<p>A place to discuss PyTorch code, issues, install, research</p>
</a>
<a class="nav-dropdown-item" href="https://pytorch.org/hub">
<span class="dropdown-title">Models (Beta)</span>
<p>Discover, publish, and reuse pre-trained models</p>
</a>
</div>
</div>
</li>
<li>
<a href="https://github.com/pytorch/pytorch">GitHub</a>
</li>
</ul>
</div>
<a class="main-menu-open-button" href="#" data-behavior="open-mobile-menu"></a>
</div>
</div>
</div>
<body class="pytorch-body">
<div class="table-of-contents-link-wrapper">
<span>Table of Contents</span>
<a href="#" class="toggle-table-of-contents" data-behavior="toggle-table-of-contents"></a>
</div>
<nav data-toggle="wy-nav-shift" class="pytorch-left-menu" id="pytorch-left-menu">
<div class="pytorch-side-scroll">
<div class="pytorch-menu pytorch-menu-vertical" data-spy="affix" role="navigation" aria-label="main navigation">
<div class="pytorch-left-menu-search">
<div class="version">
<a href='https://pytorch.org/docs/versions.html'>1.11.0 ▼</a>
</div>
<div role="search">
<form id="rtd-search-form" class="wy-form" action="search.html" method="get">
<input type="text" name="q" placeholder="Search Docs" />
<input type="hidden" name="check_keywords" value="yes" />
<input type="hidden" name="area" value="default" />
</form>
</div>
</div>
<p class="caption"><span class="caption-text">Notes</span></p>
<ul>
<li class="toctree-l1"><a class="reference internal" href="notes/amp_examples.html">Automatic Mixed Precision examples</a></li>
<li class="toctree-l1"><a class="reference internal" href="notes/autograd.html">Autograd mechanics</a></li>
<li class="toctree-l1"><a class="reference internal" href="notes/broadcasting.html">Broadcasting semantics</a></li>
<li class="toctree-l1"><a class="reference internal" href="notes/cpu_threading_torchscript_inference.html">CPU threading and TorchScript inference</a></li>
<li class="toctree-l1"><a class="reference internal" href="notes/cuda.html">CUDA semantics</a></li>
<li class="toctree-l1"><a class="reference internal" href="notes/ddp.html">Distributed Data Parallel</a></li>
<li class="toctree-l1"><a class="reference internal" href="notes/extending.html">Extending PyTorch</a></li>
<li class="toctree-l1"><a class="reference internal" href="notes/faq.html">Frequently Asked Questions</a></li>
<li class="toctree-l1"><a class="reference internal" href="notes/gradcheck.html">Gradcheck mechanics</a></li>
<li class="toctree-l1"><a class="reference internal" href="notes/hip.html">HIP (ROCm) semantics</a></li>
<li class="toctree-l1"><a class="reference internal" href="notes/large_scale_deployments.html">Features for large-scale deployments</a></li>
<li class="toctree-l1"><a class="reference internal" href="notes/modules.html">Modules</a></li>
<li class="toctree-l1"><a class="reference internal" href="notes/multiprocessing.html">Multiprocessing best practices</a></li>
<li class="toctree-l1"><a class="reference internal" href="notes/numerical_accuracy.html">Numerical accuracy</a></li>
<li class="toctree-l1"><a class="reference internal" href="notes/randomness.html">Reproducibility</a></li>
<li class="toctree-l1"><a class="reference internal" href="notes/serialization.html">Serialization semantics</a></li>
<li class="toctree-l1"><a class="reference internal" href="notes/windows.html">Windows FAQ</a></li>
</ul>
<p class="caption"><span class="caption-text">Language Bindings</span></p>
<ul>
<li class="toctree-l1"><a class="reference internal" href="cpp_index.html">C++</a></li>
<li class="toctree-l1"><a class="reference external" href="https://pytorch.org/javadoc/">Javadoc</a></li>
<li class="toctree-l1"><a class="reference internal" href="deploy.html">torch::deploy</a></li>
</ul>
<p class="caption"><span class="caption-text">Python API</span></p>
<ul class="current">
<li class="toctree-l1"><a class="reference internal" href="torch.html">torch</a></li>
<li class="toctree-l1"><a class="reference internal" href="nn.html">torch.nn</a></li>
<li class="toctree-l1"><a class="reference internal" href="nn.functional.html">torch.nn.functional</a></li>
<li class="toctree-l1"><a class="reference internal" href="tensors.html">torch.Tensor</a></li>
<li class="toctree-l1"><a class="reference internal" href="tensor_attributes.html">Tensor Attributes</a></li>
<li class="toctree-l1"><a class="reference internal" href="tensor_view.html">Tensor Views</a></li>
<li class="toctree-l1"><a class="reference internal" href="autograd.html">torch.autograd</a></li>
<li class="toctree-l1"><a class="reference internal" href="cuda.html">torch.cuda</a></li>
<li class="toctree-l1"><a class="reference internal" href="amp.html">torch.cuda.amp</a></li>
<li class="toctree-l1"><a class="reference internal" href="backends.html">torch.backends</a></li>
<li class="toctree-l1"><a class="reference internal" href="distributed.html">torch.distributed</a></li>
<li class="toctree-l1"><a class="reference internal" href="distributed.algorithms.join.html">torch.distributed.algorithms.join</a></li>
<li class="toctree-l1"><a class="reference internal" href="distributed.elastic.html">torch.distributed.elastic</a></li>
<li class="toctree-l1"><a class="reference internal" href="fsdp.html">torch.distributed.fsdp</a></li>
<li class="toctree-l1"><a class="reference internal" href="distributed.optim.html">torch.distributed.optim</a></li>
<li class="toctree-l1"><a class="reference internal" href="distributions.html">torch.distributions</a></li>
<li class="toctree-l1"><a class="reference internal" href="fft.html">torch.fft</a></li>
<li class="toctree-l1"><a class="reference internal" href="futures.html">torch.futures</a></li>
<li class="toctree-l1"><a class="reference internal" href="fx.html">torch.fx</a></li>
<li class="toctree-l1"><a class="reference internal" href="hub.html">torch.hub</a></li>
<li class="toctree-l1 current"><a class="reference internal" href="jit.html">torch.jit</a></li>
<li class="toctree-l1"><a class="reference internal" href="linalg.html">torch.linalg</a></li>
<li class="toctree-l1"><a class="reference internal" href="monitor.html">torch.monitor</a></li>
<li class="toctree-l1"><a class="reference internal" href="special.html">torch.special</a></li>
<li class="toctree-l1"><a class="reference internal" href="torch.overrides.html">torch.overrides</a></li>
<li class="toctree-l1"><a class="reference internal" href="package.html">torch.package</a></li>
<li class="toctree-l1"><a class="reference internal" href="profiler.html">torch.profiler</a></li>
<li class="toctree-l1"><a class="reference internal" href="nn.init.html">torch.nn.init</a></li>
<li class="toctree-l1"><a class="reference internal" href="onnx.html">torch.onnx</a></li>
<li class="toctree-l1"><a class="reference internal" href="optim.html">torch.optim</a></li>
<li class="toctree-l1"><a class="reference internal" href="complex_numbers.html">Complex Numbers</a></li>
<li class="toctree-l1"><a class="reference internal" href="ddp_comm_hooks.html">DDP Communication Hooks</a></li>
<li class="toctree-l1"><a class="reference internal" href="pipeline.html">Pipeline Parallelism</a></li>
<li class="toctree-l1"><a class="reference internal" href="quantization.html">Quantization</a></li>
<li class="toctree-l1"><a class="reference internal" href="rpc.html">Distributed RPC Framework</a></li>
<li class="toctree-l1"><a class="reference internal" href="random.html">torch.random</a></li>
<li class="toctree-l1"><a class="reference internal" href="sparse.html">torch.sparse</a></li>
<li class="toctree-l1"><a class="reference internal" href="storage.html">torch.Storage</a></li>
<li class="toctree-l1"><a class="reference internal" href="testing.html">torch.testing</a></li>
<li class="toctree-l1"><a class="reference internal" href="benchmark_utils.html">torch.utils.benchmark</a></li>
<li class="toctree-l1"><a class="reference internal" href="bottleneck.html">torch.utils.bottleneck</a></li>
<li class="toctree-l1"><a class="reference internal" href="checkpoint.html">torch.utils.checkpoint</a></li>
<li class="toctree-l1"><a class="reference internal" href="cpp_extension.html">torch.utils.cpp_extension</a></li>
<li class="toctree-l1"><a class="reference internal" href="data.html">torch.utils.data</a></li>
<li class="toctree-l1"><a class="reference internal" href="dlpack.html">torch.utils.dlpack</a></li>
<li class="toctree-l1"><a class="reference internal" href="mobile_optimizer.html">torch.utils.mobile_optimizer</a></li>
<li class="toctree-l1"><a class="reference internal" href="model_zoo.html">torch.utils.model_zoo</a></li>
<li class="toctree-l1"><a class="reference internal" href="tensorboard.html">torch.utils.tensorboard</a></li>
<li class="toctree-l1"><a class="reference internal" href="type_info.html">Type Info</a></li>
<li class="toctree-l1"><a class="reference internal" href="named_tensor.html">Named Tensors</a></li>
<li class="toctree-l1"><a class="reference internal" href="name_inference.html">Named Tensors operator coverage</a></li>
<li class="toctree-l1"><a class="reference internal" href="__config__.html">torch.__config__</a></li>
</ul>
<p class="caption"><span class="caption-text">Libraries</span></p>
<ul>
<li class="toctree-l1"><a class="reference external" href="https://pytorch.org/audio/stable">torchaudio</a></li>
<li class="toctree-l1"><a class="reference external" href="https://pytorch.org/text/stable">torchtext</a></li>
<li class="toctree-l1"><a class="reference external" href="https://pytorch.org/vision/stable">torchvision</a></li>
<li class="toctree-l1"><a class="reference external" href="https://pytorch.org/serve">TorchServe</a></li>
<li class="toctree-l1"><a class="reference external" href="http://pytorch.org/xla/">PyTorch on XLA Devices</a></li>
</ul>
<p class="caption"><span class="caption-text">Community</span></p>
<ul>
<li class="toctree-l1"><a class="reference internal" href="community/contribution_guide.html">PyTorch Contribution Guide</a></li>
<li class="toctree-l1"><a class="reference internal" href="community/governance.html">PyTorch Governance</a></li>
<li class="toctree-l1"><a class="reference internal" href="community/persons_of_interest.html">PyTorch Governance | Persons of Interest</a></li>
</ul>
</div>
</div>
</nav>
<div class="pytorch-container">
<div class="pytorch-page-level-bar" id="pytorch-page-level-bar">
<div class="pytorch-breadcrumbs-wrapper">
<div role="navigation" aria-label="breadcrumbs navigation">
<ul class="pytorch-breadcrumbs">
<li>
<a href="index.html">
Docs
</a> >
</li>
<li><a href="jit.html">TorchScript</a> ></li>
<li>TorchScript Language Reference</li>
<li class="pytorch-breadcrumbs-aside">
<a href="_sources/jit_language_reference_v2.rst.txt" rel="nofollow"><img src="_static/images/view-page-source-icon.svg"></a>
</li>
</ul>
</div>
</div>
<div class="pytorch-shortcuts-wrapper" id="pytorch-shortcuts-wrapper">
Shortcuts
</div>
</div>
<section data-toggle="wy-nav-shift" id="pytorch-content-wrap" class="pytorch-content-wrap">
<div class="pytorch-content-left">
<div class="rst-content">
<div role="main" class="main-content" itemscope="itemscope" itemtype="http://schema.org/Article">
<article itemprop="articleBody" id="pytorch-article" class="pytorch-article">
<div class="section" id="torchscript-language-reference">
<span id="language-reference-v2"></span><h1>TorchScript Language Reference<a class="headerlink" href="#torchscript-language-reference" title="Permalink to this headline">¶</a></h1>
<p>This reference manual describes the syntax and core semantics of the TorchScript language.
TorchScript is a statically typed subset of the Python language. This document explains the supported features of
Python in TorchScript and also how the language diverges from regular Python. Any features of Python that are not mentioned in
this reference manual are not part of TorchScript. TorchScript focuses specifically on the features of Python that are needed to
represent neural network models in PyTorch.</p>
<div class="contents local topic" id="contents">
<ul class="simple">
<li><p><a class="reference internal" href="#terminology" id="id7">Terminology</a></p></li>
<li><p><a class="reference internal" href="#id1" id="id8">Type System</a></p></li>
<li><p><a class="reference internal" href="#type-annotation" id="id9">Type Annotation</a></p></li>
<li><p><a class="reference internal" href="#expressions" id="id10">Expressions</a></p></li>
<li><p><a class="reference internal" href="#simple-statements" id="id11">Simple Statements</a></p></li>
<li><p><a class="reference internal" href="#compound-statements" id="id12">Compound Statements</a></p></li>
<li><p><a class="reference internal" href="#python-values" id="id13">Python Values</a></p></li>
<li><p><a class="reference internal" href="#torch-apis" id="id14">torch.* APIs</a></p></li>
</ul>
</div>
<div class="section" id="terminology">
<span id="type-system"></span><h2><a class="toc-backref" href="#id7">Terminology</a><a class="headerlink" href="#terminology" title="Permalink to this headline">¶</a></h2>
<p>This document uses the following terminologies:</p>
<table class="colwidths-given docutils colwidths-auto align-default">
<colgroup>
<col style="width: 50%" />
<col style="width: 50%" />
</colgroup>
<thead>
<tr class="row-odd"><th class="head"><p>Pattern</p></th>
<th class="head"><p>Notes</p></th>
</tr>
</thead>
<tbody>
<tr class="row-even"><td><p><code class="docutils literal notranslate"><span class="pre">::=</span></code></p></td>
<td><p>Indicates that the given symbol is defined as.</p></td>
</tr>
<tr class="row-odd"><td><p><code class="docutils literal notranslate"><span class="pre">"</span> <span class="pre">"</span></code></p></td>
<td><p>Represents real keywords and delimiters that are part of the syntax.</p></td>
</tr>
<tr class="row-even"><td><p><code class="docutils literal notranslate"><span class="pre">A</span> <span class="pre">|</span> <span class="pre">B</span></code></p></td>
<td><p>Indicates either A or B.</p></td>
</tr>
<tr class="row-odd"><td><p><code class="docutils literal notranslate"><span class="pre">(</span> <span class="pre">)</span></code></p></td>
<td><p>Indicates grouping.</p></td>
</tr>
<tr class="row-even"><td><p><code class="docutils literal notranslate"><span class="pre">[]</span></code></p></td>
<td><p>Indicates optional.</p></td>
</tr>
<tr class="row-odd"><td><p><code class="docutils literal notranslate"><span class="pre">A+</span></code></p></td>
<td><p>Indicates a regular expression where term A is repeated at least once.</p></td>
</tr>
<tr class="row-even"><td><p><code class="docutils literal notranslate"><span class="pre">A*</span></code></p></td>
<td><p>Indicates a regular expression where term A is repeated zero or more times.</p></td>
</tr>
</tbody>
</table>
</div>
<div class="section" id="id1">
<h2><a class="toc-backref" href="#id8">Type System</a><a class="headerlink" href="#id1" title="Permalink to this headline">¶</a></h2>
<p>TorchScript is a statically typed subset of Python. The largest difference between TorchScript and the full Python language is that TorchScript only supports a small set of types that are needed to express
neural net models.</p>
<div class="section" id="torchscript-types">
<h3>TorchScript Types<a class="headerlink" href="#torchscript-types" title="Permalink to this headline">¶</a></h3>
<p>The TorchScript type system consists of <code class="docutils literal notranslate"><span class="pre">TSType</span></code> and <code class="docutils literal notranslate"><span class="pre">TSModuleType</span></code> as defined below.</p>
<div class="highlight-default notranslate"><div class="highlight"><pre><span></span><span class="n">TSAllType</span> <span class="p">:</span><span class="o">:=</span> <span class="n">TSType</span> <span class="o">|</span> <span class="n">TSModuleType</span>
<span class="n">TSType</span> <span class="p">:</span><span class="o">:=</span> <span class="n">TSMetaType</span> <span class="o">|</span> <span class="n">TSPrimitiveType</span> <span class="o">|</span> <span class="n">TSStructuralType</span> <span class="o">|</span> <span class="n">TSNominalType</span>
</pre></div>
</div>
<p><code class="docutils literal notranslate"><span class="pre">TSType</span></code> represents the majority of TorchScript types that are composable and that can be used in TorchScript type annotations.
<code class="docutils literal notranslate"><span class="pre">TSType</span></code> refers to any of the following:</p>
<ul class="simple">
<li><p>Meta Types, e.g., <code class="docutils literal notranslate"><span class="pre">Any</span></code></p></li>
<li><p>Primitive Types, e.g., <code class="docutils literal notranslate"><span class="pre">int</span></code>, <code class="docutils literal notranslate"><span class="pre">float</span></code>, and <code class="docutils literal notranslate"><span class="pre">str</span></code></p></li>
<li><p>Structural Types, e.g., <code class="docutils literal notranslate"><span class="pre">Optional[int]</span></code> or <code class="docutils literal notranslate"><span class="pre">List[MyClass]</span></code></p></li>
<li><p>Nominal Types (Python classes), e.g., <code class="docutils literal notranslate"><span class="pre">MyClass</span></code> (user-defined), <code class="docutils literal notranslate"><span class="pre">torch.tensor</span></code> (built-in)</p></li>
</ul>
<p><code class="docutils literal notranslate"><span class="pre">TSModuleType</span></code> represents <code class="docutils literal notranslate"><span class="pre">torch.nn.Module</span></code> and its subclasses. It is treated differently from <code class="docutils literal notranslate"><span class="pre">TSType</span></code> because its type schema is inferred partly from the object instance and partly from the class definition.
As such, instances of a <code class="docutils literal notranslate"><span class="pre">TSModuleType</span></code> may not follow the same static type schema. <code class="docutils literal notranslate"><span class="pre">TSModuleType</span></code> cannot be used as a TorchScript type annotation or be composed with <code class="docutils literal notranslate"><span class="pre">TSType</span></code> for type safety considerations.</p>
</div>
<div class="section" id="meta-types">
<h3>Meta Types<a class="headerlink" href="#meta-types" title="Permalink to this headline">¶</a></h3>
<p>Meta types are so abstract that they are more like type constraints than concrete types.
Currently TorchScript defines one meta-type, <code class="docutils literal notranslate"><span class="pre">Any</span></code>, that represents any TorchScript type.</p>
<div class="section" id="any-type">
<h4><code class="docutils literal notranslate"><span class="pre">Any</span></code> Type<a class="headerlink" href="#any-type" title="Permalink to this headline">¶</a></h4>
<p>The <code class="docutils literal notranslate"><span class="pre">Any</span></code> type represents any TorchScript type. <code class="docutils literal notranslate"><span class="pre">Any</span></code> specifies no type constraints, thus there is no type-checking on <code class="docutils literal notranslate"><span class="pre">Any</span></code>.
As such it can be bound to any Python or TorchScript data types (e.g., <code class="docutils literal notranslate"><span class="pre">int</span></code>, TorchScript <code class="docutils literal notranslate"><span class="pre">tuple</span></code>, or an arbitrary Python class that is not scripted).</p>
<div class="highlight-default notranslate"><div class="highlight"><pre><span></span><span class="n">TSMetaType</span> <span class="p">:</span><span class="o">:=</span> <span class="s2">"Any"</span>
</pre></div>
</div>
<p>Where:</p>
<ul class="simple">
<li><p><code class="docutils literal notranslate"><span class="pre">Any</span></code> is the Python class name from the typing module. Therefore, to use the <code class="docutils literal notranslate"><span class="pre">Any</span></code> type, you must import it from <code class="docutils literal notranslate"><span class="pre">typing</span></code> (e.g., <code class="docutils literal notranslate"><span class="pre">from</span> <span class="pre">typing</span> <span class="pre">import</span> <span class="pre">Any</span></code>).</p></li>
<li><p>Since <code class="docutils literal notranslate"><span class="pre">Any</span></code> can represent any TorchScript type, the set of operators that are allowed to operate on values of this type on <code class="docutils literal notranslate"><span class="pre">Any</span></code> is limited.</p></li>
</ul>
</div>
<div class="section" id="operators-supported-for-any-type">
<h4>Operators Supported for <code class="docutils literal notranslate"><span class="pre">Any</span></code> Type<a class="headerlink" href="#operators-supported-for-any-type" title="Permalink to this headline">¶</a></h4>
<ul class="simple">
<li><p>Assignment to data of <code class="docutils literal notranslate"><span class="pre">Any</span></code> type.</p></li>
<li><p>Binding to parameter or return of <code class="docutils literal notranslate"><span class="pre">Any</span></code> type.</p></li>
<li><p><code class="docutils literal notranslate"><span class="pre">x</span> <span class="pre">is</span></code>, <code class="docutils literal notranslate"><span class="pre">x</span> <span class="pre">is</span> <span class="pre">not</span></code> where <code class="docutils literal notranslate"><span class="pre">x</span></code> is of <code class="docutils literal notranslate"><span class="pre">Any</span></code> type.</p></li>
<li><p><code class="docutils literal notranslate"><span class="pre">isinstance(x,</span> <span class="pre">Type)</span></code> where <code class="docutils literal notranslate"><span class="pre">x</span></code> is of <code class="docutils literal notranslate"><span class="pre">Any</span></code> type.</p></li>
<li><p>Data of <code class="docutils literal notranslate"><span class="pre">Any</span></code> type is printable.</p></li>
<li><p>Data of <code class="docutils literal notranslate"><span class="pre">List[Any]</span></code> type may be sortable if the data is a list of values of the same type <code class="docutils literal notranslate"><span class="pre">T</span></code> and that <code class="docutils literal notranslate"><span class="pre">T</span></code> supports comparison operators.</p></li>
</ul>
<p><strong>Compared to Python</strong></p>
<p><code class="docutils literal notranslate"><span class="pre">Any</span></code> is the least constrained type in the TorchScript type system. In that sense, it is quite similar to the
<code class="docutils literal notranslate"><span class="pre">Object</span></code> class in Python. However, <code class="docutils literal notranslate"><span class="pre">Any</span></code> only supports a subset of the operators and methods that are supported by <code class="docutils literal notranslate"><span class="pre">Object</span></code>.</p>
</div>
<div class="section" id="design-notes">
<h4>Design Notes<a class="headerlink" href="#design-notes" title="Permalink to this headline">¶</a></h4>
<p>When we script a PyTorch module, we may encounter data that is not involved in the execution of the script. Nevertheless, it has to be described
by a type schema. It is not only cumbersome to describe static types for unused data (in the context of the script), but also may lead to unnecessary
scripting failures. <code class="docutils literal notranslate"><span class="pre">Any</span></code> is introduced to describe the type of the data where precise static types are not necessary for compilation.</p>
<p><strong>Example 1</strong></p>
<p>This example illustrates how <code class="docutils literal notranslate"><span class="pre">Any</span></code> can be used to allow the second element of the tuple parameter to be of any type. This is possible
because <code class="docutils literal notranslate"><span class="pre">x[1]</span></code> is not involved in any computation that requires knowing its precise type.</p>
<div class="highlight-python3 notranslate"><div class="highlight"><pre><span></span><span class="kn">import</span> <span class="nn">torch</span>
<span class="kn">from</span> <span class="nn">typing</span> <span class="kn">import</span> <span class="n">Tuple</span>
<span class="kn">from</span> <span class="nn">typing</span> <span class="kn">import</span> <span class="n">Any</span>
<span class="nd">@torch</span><span class="o">.</span><span class="n">jit</span><span class="o">.</span><span class="n">export</span>
<span class="k">def</span> <span class="nf">inc_first_element</span><span class="p">(</span><span class="n">x</span><span class="p">:</span> <span class="n">Tuple</span><span class="p">[</span><span class="nb">int</span><span class="p">,</span> <span class="n">Any</span><span class="p">]):</span>
<span class="k">return</span> <span class="p">(</span><span class="n">x</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span><span class="o">+</span><span class="mi">1</span><span class="p">,</span> <span class="n">x</span><span class="p">[</span><span class="mi">1</span><span class="p">])</span>
<span class="n">m</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">jit</span><span class="o">.</span><span class="n">script</span><span class="p">(</span><span class="n">inc_first_element</span><span class="p">)</span>
<span class="nb">print</span><span class="p">(</span><span class="n">m</span><span class="p">((</span><span class="mi">1</span><span class="p">,</span><span class="mf">2.0</span><span class="p">)))</span>
<span class="nb">print</span><span class="p">(</span><span class="n">m</span><span class="p">((</span><span class="mi">1</span><span class="p">,(</span><span class="mi">100</span><span class="p">,</span><span class="mi">200</span><span class="p">))))</span>
</pre></div>
</div>
<p>The example above produces the following output:</p>
<div class="highlight-none notranslate"><div class="highlight"><pre><span></span>(2, 2.0)
(2, (100, 200))
</pre></div>
</div>
<p>The second element of the tuple is of <code class="docutils literal notranslate"><span class="pre">Any</span></code> type, thus can bind to multiple types.
For example, <code class="docutils literal notranslate"><span class="pre">(1,</span> <span class="pre">2.0)</span></code> binds a float type to <code class="docutils literal notranslate"><span class="pre">Any</span></code> as in <code class="docutils literal notranslate"><span class="pre">Tuple[int,</span> <span class="pre">Any]</span></code>,
whereas <code class="docutils literal notranslate"><span class="pre">(1,</span> <span class="pre">(100,</span> <span class="pre">200))</span></code> binds a tuple to <code class="docutils literal notranslate"><span class="pre">Any</span></code> in the second invocation.</p>
<p><strong>Example 2</strong></p>
<p>This example illustrates how we can use <code class="docutils literal notranslate"><span class="pre">isinstance</span></code> to dynamically check the type of the data that is annotated as <code class="docutils literal notranslate"><span class="pre">Any</span></code> type:</p>
<div class="highlight-python3 notranslate"><div class="highlight"><pre><span></span><span class="kn">import</span> <span class="nn">torch</span>
<span class="kn">from</span> <span class="nn">typing</span> <span class="kn">import</span> <span class="n">Any</span>
<span class="k">def</span> <span class="nf">f</span><span class="p">(</span><span class="n">a</span><span class="p">:</span><span class="n">Any</span><span class="p">):</span>
<span class="nb">print</span><span class="p">(</span><span class="n">a</span><span class="p">)</span>
<span class="k">return</span> <span class="p">(</span><span class="nb">isinstance</span><span class="p">(</span><span class="n">a</span><span class="p">,</span> <span class="n">torch</span><span class="o">.</span><span class="n">Tensor</span><span class="p">))</span>
<span class="n">ones</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">ones</span><span class="p">([</span><span class="mi">2</span><span class="p">])</span>
<span class="n">m</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">jit</span><span class="o">.</span><span class="n">script</span><span class="p">(</span><span class="n">f</span><span class="p">)</span>
<span class="nb">print</span><span class="p">(</span><span class="n">m</span><span class="p">(</span><span class="n">ones</span><span class="p">))</span>
</pre></div>
</div>
<p>The example above produces the following output:</p>
<div class="highlight-none notranslate"><div class="highlight"><pre><span></span> 1
1
[ CPUFloatType{2} ]
True
</pre></div>
</div>
</div>
</div>
<div class="section" id="primitive-types">
<h3>Primitive Types<a class="headerlink" href="#primitive-types" title="Permalink to this headline">¶</a></h3>
<p>Primitive TorchScript types are types that represent a single type of value and go with a single pre-defined
type name.</p>
<div class="highlight-default notranslate"><div class="highlight"><pre><span></span><span class="n">TSPrimitiveType</span> <span class="p">:</span><span class="o">:=</span> <span class="s2">"int"</span> <span class="o">|</span> <span class="s2">"float"</span> <span class="o">|</span> <span class="s2">"double"</span> <span class="o">|</span> <span class="s2">"complex"</span> <span class="o">|</span> <span class="s2">"bool"</span> <span class="o">|</span> <span class="s2">"str"</span> <span class="o">|</span> <span class="s2">"None"</span>
</pre></div>
</div>
</div>
<div class="section" id="structural-types">
<h3>Structural Types<a class="headerlink" href="#structural-types" title="Permalink to this headline">¶</a></h3>
<p>Structural types are types that are structurally defined without a user-defined name (unlike nominal types),
such as <code class="docutils literal notranslate"><span class="pre">Future[int]</span></code>. Structural types are composable with any <code class="docutils literal notranslate"><span class="pre">TSType</span></code>.</p>
<div class="highlight-default notranslate"><div class="highlight"><pre><span></span><span class="n">TSStructuralType</span> <span class="p">:</span><span class="o">:=</span> <span class="n">TSTuple</span> <span class="o">|</span> <span class="n">TSNamedTuple</span> <span class="o">|</span> <span class="n">TSList</span> <span class="o">|</span> <span class="n">TSDict</span> <span class="o">|</span>
<span class="n">TSOptional</span> <span class="o">|</span> <span class="n">TSUnion</span> <span class="o">|</span> <span class="n">TSFuture</span> <span class="o">|</span> <span class="n">TSRRef</span>
<span class="n">TSTuple</span> <span class="p">:</span><span class="o">:=</span> <span class="s2">"Tuple"</span> <span class="s2">"["</span> <span class="p">(</span><span class="n">TSType</span> <span class="s2">","</span><span class="p">)</span><span class="o">*</span> <span class="n">TSType</span> <span class="s2">"]"</span>
<span class="n">TSNamedTuple</span> <span class="p">:</span><span class="o">:=</span> <span class="s2">"namedtuple"</span> <span class="s2">"("</span> <span class="p">(</span><span class="n">TSType</span> <span class="s2">","</span><span class="p">)</span><span class="o">*</span> <span class="n">TSType</span> <span class="s2">")"</span>
<span class="n">TSList</span> <span class="p">:</span><span class="o">:=</span> <span class="s2">"List"</span> <span class="s2">"["</span> <span class="n">TSType</span> <span class="s2">"]"</span>
<span class="n">TSOptional</span> <span class="p">:</span><span class="o">:=</span> <span class="s2">"Optional"</span> <span class="s2">"["</span> <span class="n">TSType</span> <span class="s2">"]"</span>
<span class="n">TSUnion</span> <span class="p">:</span><span class="o">:=</span> <span class="s2">"Union"</span> <span class="s2">"["</span> <span class="p">(</span><span class="n">TSType</span> <span class="s2">","</span><span class="p">)</span><span class="o">*</span> <span class="n">TSType</span> <span class="s2">"]"</span>
<span class="n">TSFuture</span> <span class="p">:</span><span class="o">:=</span> <span class="s2">"Future"</span> <span class="s2">"["</span> <span class="n">TSType</span> <span class="s2">"]"</span>
<span class="n">TSRRef</span> <span class="p">:</span><span class="o">:=</span> <span class="s2">"RRef"</span> <span class="s2">"["</span> <span class="n">TSType</span> <span class="s2">"]"</span>
<span class="n">TSDict</span> <span class="p">:</span><span class="o">:=</span> <span class="s2">"Dict"</span> <span class="s2">"["</span> <span class="n">KeyType</span> <span class="s2">","</span> <span class="n">TSType</span> <span class="s2">"]"</span>
<span class="n">KeyType</span> <span class="p">:</span><span class="o">:=</span> <span class="s2">"str"</span> <span class="o">|</span> <span class="s2">"int"</span> <span class="o">|</span> <span class="s2">"float"</span> <span class="o">|</span> <span class="s2">"bool"</span> <span class="o">|</span> <span class="n">TensorType</span> <span class="o">|</span> <span class="s2">"Any"</span>
</pre></div>
</div>
<p>Where:</p>
<ul class="simple">
<li><p><code class="docutils literal notranslate"><span class="pre">Tuple</span></code>, <code class="docutils literal notranslate"><span class="pre">List</span></code>, <code class="docutils literal notranslate"><span class="pre">Optional</span></code>, <code class="docutils literal notranslate"><span class="pre">Union</span></code>, <code class="docutils literal notranslate"><span class="pre">Future</span></code>, <code class="docutils literal notranslate"><span class="pre">Dict</span></code> represent Python type class names that are defined in the module <code class="docutils literal notranslate"><span class="pre">typing</span></code>. To use these type names, you must import them from <code class="docutils literal notranslate"><span class="pre">typing</span></code> (e.g., <code class="docutils literal notranslate"><span class="pre">from</span> <span class="pre">typing</span> <span class="pre">import</span> <span class="pre">Tuple</span></code>).</p></li>
<li><p><code class="docutils literal notranslate"><span class="pre">namedtuple</span></code> represents the Python class <code class="docutils literal notranslate"><span class="pre">collections.namedtuple</span></code> or <code class="docutils literal notranslate"><span class="pre">typing.NamedTuple</span></code>.</p></li>
<li><p><code class="docutils literal notranslate"><span class="pre">Future</span></code> and <code class="docutils literal notranslate"><span class="pre">RRef</span></code> represent the Python classes <code class="docutils literal notranslate"><span class="pre">torch.futures</span></code> and <code class="docutils literal notranslate"><span class="pre">torch.distributed.rpc</span></code>.</p></li>
</ul>
<p><strong>Compared to Python</strong></p>
<p>Apart from being composable with TorchScript types, these TorchScript structural types often support a common subset of the operators and methods of their Python counterparts.</p>
<p><strong>Example 1</strong></p>
<p>This example uses <code class="docutils literal notranslate"><span class="pre">typing.NamedTuple</span></code> syntax to define a tuple:</p>
<div class="highlight-python3 notranslate"><div class="highlight"><pre><span></span><span class="kn">import</span> <span class="nn">torch</span>
<span class="kn">from</span> <span class="nn">typing</span> <span class="kn">import</span> <span class="n">NamedTuple</span>
<span class="kn">from</span> <span class="nn">typing</span> <span class="kn">import</span> <span class="n">Tuple</span>
<span class="k">class</span> <span class="nc">MyTuple</span><span class="p">(</span><span class="n">NamedTuple</span><span class="p">):</span>
<span class="n">first</span><span class="p">:</span> <span class="nb">int</span>
<span class="n">second</span><span class="p">:</span> <span class="nb">int</span>
<span class="k">def</span> <span class="nf">inc</span><span class="p">(</span><span class="n">x</span><span class="p">:</span> <span class="n">MyTuple</span><span class="p">)</span> <span class="o">-></span> <span class="n">Tuple</span><span class="p">[</span><span class="nb">int</span><span class="p">,</span> <span class="nb">int</span><span class="p">]:</span>
<span class="k">return</span> <span class="p">(</span><span class="n">x</span><span class="o">.</span><span class="n">first</span><span class="o">+</span><span class="mi">1</span><span class="p">,</span> <span class="n">x</span><span class="o">.</span><span class="n">second</span><span class="o">+</span><span class="mi">1</span><span class="p">)</span>
<span class="n">t</span> <span class="o">=</span> <span class="n">MyTuple</span><span class="p">(</span><span class="n">first</span><span class="o">=</span><span class="mi">1</span><span class="p">,</span> <span class="n">second</span><span class="o">=</span><span class="mi">2</span><span class="p">)</span>
<span class="n">scripted_inc</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">jit</span><span class="o">.</span><span class="n">script</span><span class="p">(</span><span class="n">inc</span><span class="p">)</span>
<span class="nb">print</span><span class="p">(</span><span class="s2">"TorchScript:"</span><span class="p">,</span> <span class="n">scripted_inc</span><span class="p">(</span><span class="n">t</span><span class="p">))</span>
</pre></div>
</div>
<p>The example above produces the following output:</p>
<div class="highlight-none notranslate"><div class="highlight"><pre><span></span>TorchScript: (2, 3)
</pre></div>
</div>
<p><strong>Example 2</strong></p>
<p>This example uses <code class="docutils literal notranslate"><span class="pre">collections.namedtuple</span></code> syntax to define a tuple:</p>
<div class="highlight-python3 notranslate"><div class="highlight"><pre><span></span><span class="kn">import</span> <span class="nn">torch</span>
<span class="kn">from</span> <span class="nn">typing</span> <span class="kn">import</span> <span class="n">NamedTuple</span>
<span class="kn">from</span> <span class="nn">typing</span> <span class="kn">import</span> <span class="n">Tuple</span>
<span class="kn">from</span> <span class="nn">collections</span> <span class="kn">import</span> <span class="n">namedtuple</span>
<span class="n">_AnnotatedNamedTuple</span> <span class="o">=</span> <span class="n">NamedTuple</span><span class="p">(</span><span class="s1">'_NamedTupleAnnotated'</span><span class="p">,</span> <span class="p">[(</span><span class="s1">'first'</span><span class="p">,</span> <span class="nb">int</span><span class="p">),</span> <span class="p">(</span><span class="s1">'second'</span><span class="p">,</span> <span class="nb">int</span><span class="p">)])</span>
<span class="n">_UnannotatedNamedTuple</span> <span class="o">=</span> <span class="n">namedtuple</span><span class="p">(</span><span class="s1">'_NamedTupleAnnotated'</span><span class="p">,</span> <span class="p">[</span><span class="s1">'first'</span><span class="p">,</span> <span class="s1">'second'</span><span class="p">])</span>
<span class="k">def</span> <span class="nf">inc</span><span class="p">(</span><span class="n">x</span><span class="p">:</span> <span class="n">_AnnotatedNamedTuple</span><span class="p">)</span> <span class="o">-></span> <span class="n">Tuple</span><span class="p">[</span><span class="nb">int</span><span class="p">,</span> <span class="nb">int</span><span class="p">]:</span>
<span class="k">return</span> <span class="p">(</span><span class="n">x</span><span class="o">.</span><span class="n">first</span><span class="o">+</span><span class="mi">1</span><span class="p">,</span> <span class="n">x</span><span class="o">.</span><span class="n">second</span><span class="o">+</span><span class="mi">1</span><span class="p">)</span>
<span class="n">m</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">jit</span><span class="o">.</span><span class="n">script</span><span class="p">(</span><span class="n">inc</span><span class="p">)</span>
<span class="nb">print</span><span class="p">(</span><span class="n">inc</span><span class="p">(</span><span class="n">_UnannotatedNamedTuple</span><span class="p">(</span><span class="mi">1</span><span class="p">,</span><span class="mi">2</span><span class="p">)))</span>
</pre></div>
</div>
<p>The example above produces the following output:</p>
<div class="highlight-none notranslate"><div class="highlight"><pre><span></span>(2, 3)
</pre></div>
</div>
<p><strong>Example 3</strong></p>
<p>This example illustrates a common mistake of annotating structural types, i.e., not importing the composite type
classes from the <code class="docutils literal notranslate"><span class="pre">typing</span></code> module:</p>
<div class="highlight-default notranslate"><div class="highlight"><pre><span></span><span class="kn">import</span> <span class="nn">torch</span>
<span class="c1"># ERROR: Tuple not recognized because not imported from typing</span>
<span class="nd">@torch</span><span class="o">.</span><span class="n">jit</span><span class="o">.</span><span class="n">export</span>
<span class="k">def</span> <span class="nf">inc</span><span class="p">(</span><span class="n">x</span><span class="p">:</span> <span class="n">Tuple</span><span class="p">[</span><span class="nb">int</span><span class="p">,</span> <span class="nb">int</span><span class="p">]):</span>
<span class="k">return</span> <span class="p">(</span><span class="n">x</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span><span class="o">+</span><span class="mi">1</span><span class="p">,</span> <span class="n">x</span><span class="p">[</span><span class="mi">1</span><span class="p">]</span><span class="o">+</span><span class="mi">1</span><span class="p">)</span>
<span class="n">m</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">jit</span><span class="o">.</span><span class="n">script</span><span class="p">(</span><span class="n">inc</span><span class="p">)</span>
<span class="nb">print</span><span class="p">(</span><span class="n">m</span><span class="p">((</span><span class="mi">1</span><span class="p">,</span><span class="mi">2</span><span class="p">)))</span>
</pre></div>
</div>
<p>Running the above code yields the following scripting error:</p>
<div class="highlight-default notranslate"><div class="highlight"><pre><span></span><span class="n">File</span> <span class="s2">"test-tuple.py"</span><span class="p">,</span> <span class="n">line</span> <span class="mi">5</span><span class="p">,</span> <span class="ow">in</span> <span class="o"><</span><span class="n">module</span><span class="o">></span>
<span class="k">def</span> <span class="nf">inc</span><span class="p">(</span><span class="n">x</span><span class="p">:</span> <span class="n">Tuple</span><span class="p">[</span><span class="nb">int</span><span class="p">,</span> <span class="nb">int</span><span class="p">]):</span>
<span class="ne">NameError</span><span class="p">:</span> <span class="n">name</span> <span class="s1">'Tuple'</span> <span class="ow">is</span> <span class="ow">not</span> <span class="n">defined</span>
</pre></div>
</div>
<p>The remedy is to add the line <code class="docutils literal notranslate"><span class="pre">from</span> <span class="pre">typing</span> <span class="pre">import</span> <span class="pre">Tuple</span></code> to the beginning of the code.</p>
</div>
<div class="section" id="nominal-types">
<h3>Nominal Types<a class="headerlink" href="#nominal-types" title="Permalink to this headline">¶</a></h3>
<p>Nominal TorchScript types are Python classes. These types are called nominal because they are declared with a custom
name and are compared using class names. Nominal classes are further classified into the following categories:</p>
<div class="highlight-default notranslate"><div class="highlight"><pre><span></span><span class="n">TSNominalType</span> <span class="p">:</span><span class="o">:=</span> <span class="n">TSBuiltinClasses</span> <span class="o">|</span> <span class="n">TSCustomClass</span> <span class="o">|</span> <span class="n">TSEnum</span>
</pre></div>
</div>
<p>Among them, <code class="docutils literal notranslate"><span class="pre">TSCustomClass</span></code> and <code class="docutils literal notranslate"><span class="pre">TSEnum</span></code> must be compilable to TorchScript Intermediate Representation (IR). This is enforced by the type-checker.</p>
</div>
<div class="section" id="built-in-class">
<h3>Built-in Class<a class="headerlink" href="#built-in-class" title="Permalink to this headline">¶</a></h3>
<p>Built-in nominal types are Python classes whose semantics are built into the TorchScript system (e.g., tensor types).
TorchScript defines the semantics of these built-in nominal types, and often supports only a subset of the methods or
attributes of its Python class definition.</p>
<div class="highlight-default notranslate"><div class="highlight"><pre><span></span><span class="n">TSBuiltinClass</span> <span class="p">:</span><span class="o">:=</span> <span class="n">TSTensor</span> <span class="o">|</span> <span class="s2">"torch.device"</span> <span class="o">|</span> <span class="s2">"torch.Stream"</span> <span class="o">|</span> <span class="s2">"torch.dtype"</span> <span class="o">|</span>
<span class="s2">"torch.nn.ModuleList"</span> <span class="o">|</span> <span class="s2">"torch.nn.ModuleDict"</span> <span class="o">|</span> <span class="o">...</span>
<span class="n">TSTensor</span> <span class="p">:</span><span class="o">:=</span> <span class="s2">"torch.Tensor"</span> <span class="o">|</span> <span class="s2">"common.SubTensor"</span> <span class="o">|</span> <span class="s2">"common.SubWithTorchFunction"</span> <span class="o">|</span>
<span class="s2">"torch.nn.parameter.Parameter"</span> <span class="o">|</span> <span class="ow">and</span> <span class="n">subclasses</span> <span class="n">of</span> <span class="n">torch</span><span class="o">.</span><span class="n">Tensor</span>
</pre></div>
</div>
<div class="section" id="special-note-on-torch-nn-modulelist-and-torch-nn-moduledict">
<h4>Special Note on torch.nn.ModuleList and torch.nn.ModuleDict<a class="headerlink" href="#special-note-on-torch-nn-modulelist-and-torch-nn-moduledict" title="Permalink to this headline">¶</a></h4>
<p>Although <code class="docutils literal notranslate"><span class="pre">torch.nn.ModuleList</span></code> and <code class="docutils literal notranslate"><span class="pre">torch.nn.ModuleDict</span></code> are defined as a list and dictionary in Python,
they behave more like tuples in TorchScript:</p>
<ul class="simple">
<li><p>In TorchScript, instances of <code class="docutils literal notranslate"><span class="pre">torch.nn.ModuleList</span></code> or <code class="docutils literal notranslate"><span class="pre">torch.nn.ModuleDict</span></code> are immutable.</p></li>
<li><p>Code that iterates over <code class="docutils literal notranslate"><span class="pre">torch.nn.ModuleList</span></code> or <code class="docutils literal notranslate"><span class="pre">torch.nn.ModuleDict</span></code> is completely unrolled so that elements of <code class="docutils literal notranslate"><span class="pre">torch.nn.ModuleList</span></code> or keys of <code class="docutils literal notranslate"><span class="pre">torch.nn.ModuleDict</span></code> can be of different subclasses of <code class="docutils literal notranslate"><span class="pre">torch.nn.Module</span></code>.</p></li>
</ul>
<p><strong>Example</strong></p>
<p>The following example highlights the use of a few built-in Torchscript classes (<code class="docutils literal notranslate"><span class="pre">torch.*</span></code>):</p>
<div class="highlight-default notranslate"><div class="highlight"><pre><span></span><span class="kn">import</span> <span class="nn">torch</span>
<span class="nd">@torch</span><span class="o">.</span><span class="n">jit</span><span class="o">.</span><span class="n">script</span>
<span class="k">class</span> <span class="nc">A</span><span class="p">:</span>
<span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
<span class="bp">self</span><span class="o">.</span><span class="n">x</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">rand</span><span class="p">(</span><span class="mi">3</span><span class="p">)</span>
<span class="k">def</span> <span class="nf">f</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">y</span><span class="p">:</span> <span class="n">torch</span><span class="o">.</span><span class="n">device</span><span class="p">):</span>
<span class="k">return</span> <span class="bp">self</span><span class="o">.</span><span class="n">x</span><span class="o">.</span><span class="n">to</span><span class="p">(</span><span class="n">device</span><span class="o">=</span><span class="n">y</span><span class="p">)</span>
<span class="k">def</span> <span class="nf">g</span><span class="p">():</span>
<span class="n">a</span> <span class="o">=</span> <span class="n">A</span><span class="p">()</span>
<span class="k">return</span> <span class="n">a</span><span class="o">.</span><span class="n">f</span><span class="p">(</span><span class="n">torch</span><span class="o">.</span><span class="n">device</span><span class="p">(</span><span class="s2">"cpu"</span><span class="p">))</span>
<span class="n">script_g</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">jit</span><span class="o">.</span><span class="n">script</span><span class="p">(</span><span class="n">g</span><span class="p">)</span>
<span class="nb">print</span><span class="p">(</span><span class="n">script_g</span><span class="o">.</span><span class="n">graph</span><span class="p">)</span>
</pre></div>
</div>
</div>
</div>
<div class="section" id="custom-class">
<h3>Custom Class<a class="headerlink" href="#custom-class" title="Permalink to this headline">¶</a></h3>
<p>Unlike built-in classes, semantics of custom classes are user-defined and the entire class definition must be compilable to TorchScript IR and subject to TorchScript type-checking rules.</p>
<div class="highlight-default notranslate"><div class="highlight"><pre><span></span><span class="n">TSClassDef</span> <span class="p">:</span><span class="o">:=</span> <span class="p">[</span> <span class="s2">"@torch.jit.script"</span> <span class="p">]</span>
<span class="s2">"class"</span> <span class="n">ClassName</span> <span class="p">[</span> <span class="s2">"(object)"</span> <span class="p">]</span> <span class="s2">":"</span>
<span class="n">MethodDefinition</span> <span class="o">|</span>
<span class="p">[</span> <span class="s2">"@torch.jit.ignore"</span> <span class="p">]</span> <span class="o">|</span> <span class="p">[</span> <span class="s2">"@torch.jit.unused"</span> <span class="p">]</span>
<span class="n">MethodDefinition</span>
</pre></div>
</div>
<p>Where:</p>
<ul class="simple">
<li><p>Classes must be new-style classes. Python 3 supports only new-style classes. In Python 2.x, a new-style class is specified by subclassing from the object.</p></li>
<li><p>Instance data attributes are statically typed, and instance attributes must be declared by assignments inside the <code class="docutils literal notranslate"><span class="pre">__init__()</span></code> method.</p></li>
<li><p>Method overloading is not supported (i.e., you cannot have multiple methods with the same method name).</p></li>
<li><p><code class="docutils literal notranslate"><span class="pre">MethodDefinition</span></code> must be compilable to TorchScript IR and adhere to TorchScript’s type-checking rules, (i.e., all methods must be valid TorchScript functions and class attribute definitions must be valid TorchScript statements).</p></li>
<li><p><code class="docutils literal notranslate"><span class="pre">torch.jit.ignore</span></code> and <code class="docutils literal notranslate"><span class="pre">torch.jit.unused</span></code> can be used to ignore the method or function that is not fully torchscriptable or should be ignored by the compiler.</p></li>
</ul>
<p><strong>Compared to Python</strong></p>
<p>TorchScript custom classes are quite limited compared to their Python counterpart. Torchscript custom classes:</p>
<ul class="simple">
<li><p>Do not support class attributes.</p></li>
<li><p>Do not support subclassing except for subclassing an interface type or object.</p></li>
<li><p>Do not support method overloading.</p></li>
<li><p>Must initialize all its instance attributes in <code class="docutils literal notranslate"><span class="pre">__init__()</span></code>; this is because TorchScript constructs a static schema of the class by inferring attribute types in <code class="docutils literal notranslate"><span class="pre">__init__()</span></code>.</p></li>
<li><p>Must contain only methods that satisfy TorchScript type-checking rules and are compilable to TorchScript IRs.</p></li>
</ul>
<p><strong>Example 1</strong></p>
<p>Python classes can be used in TorchScript if they are annotated with <code class="docutils literal notranslate"><span class="pre">@torch.jit.script</span></code>, similar to how a TorchScript function would be declared:</p>
<div class="highlight-default notranslate"><div class="highlight"><pre><span></span><span class="nd">@torch</span><span class="o">.</span><span class="n">jit</span><span class="o">.</span><span class="n">script</span>
<span class="k">class</span> <span class="nc">MyClass</span><span class="p">:</span>
<span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">x</span><span class="p">:</span> <span class="nb">int</span><span class="p">):</span>
<span class="bp">self</span><span class="o">.</span><span class="n">x</span> <span class="o">=</span> <span class="n">x</span>
<span class="k">def</span> <span class="nf">inc</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">val</span><span class="p">:</span> <span class="nb">int</span><span class="p">):</span>
<span class="bp">self</span><span class="o">.</span><span class="n">x</span> <span class="o">+=</span> <span class="n">val</span>
</pre></div>
</div>
<p><strong>Example 2</strong></p>
<p>A TorchScript custom class type must “declare” all its instance attributes by assignments in <code class="docutils literal notranslate"><span class="pre">__init__()</span></code>. If an instance attribute is not defined in <code class="docutils literal notranslate"><span class="pre">__init__()</span></code> but accessed in other methods of the class, the class cannot be compiled as a TorchScript class, as shown in the following example:</p>
<div class="highlight-default notranslate"><div class="highlight"><pre><span></span><span class="kn">import</span> <span class="nn">torch</span>
<span class="nd">@torch</span><span class="o">.</span><span class="n">jit</span><span class="o">.</span><span class="n">script</span>
<span class="k">class</span> <span class="nc">foo</span><span class="p">:</span>
<span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
<span class="bp">self</span><span class="o">.</span><span class="n">y</span> <span class="o">=</span> <span class="mi">1</span>
<span class="c1"># ERROR: self.x is not defined in __init__</span>
<span class="k">def</span> <span class="nf">assign_x</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
<span class="bp">self</span><span class="o">.</span><span class="n">x</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">rand</span><span class="p">(</span><span class="mi">2</span><span class="p">,</span> <span class="mi">3</span><span class="p">)</span>
</pre></div>
</div>
<p>The class will fail to compile and issue the following error:</p>
<div class="highlight-default notranslate"><div class="highlight"><pre><span></span>RuntimeError:
Tried to set nonexistent attribute: x. Did you forget to initialize it in __init__()?:
def assign_x(self):
self.x = torch.rand(2, 3)
~~~~~~~~~~~~~~~~~~~~~~~~ <--- HERE
</pre></div>
</div>
<p><strong>Example 3</strong></p>
<p>In this example, a TorchScript custom class defines a class variable name, which is not allowed:</p>
<div class="highlight-default notranslate"><div class="highlight"><pre><span></span><span class="kn">import</span> <span class="nn">torch</span>
<span class="nd">@torch</span><span class="o">.</span><span class="n">jit</span><span class="o">.</span><span class="n">script</span>
<span class="k">class</span> <span class="nc">MyClass</span><span class="p">(</span><span class="nb">object</span><span class="p">):</span>
<span class="n">name</span> <span class="o">=</span> <span class="s2">"MyClass"</span>
<span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">x</span><span class="p">:</span> <span class="nb">int</span><span class="p">):</span>
<span class="bp">self</span><span class="o">.</span><span class="n">x</span> <span class="o">=</span> <span class="n">x</span>
<span class="k">def</span> <span class="nf">fn</span><span class="p">(</span><span class="n">a</span><span class="p">:</span> <span class="n">MyClass</span><span class="p">):</span>
<span class="k">return</span> <span class="n">a</span><span class="o">.</span><span class="n">name</span>
</pre></div>
</div>
<p>It leads to the following compile-time error:</p>
<div class="highlight-default notranslate"><div class="highlight"><pre><span></span>RuntimeError:
'__torch__.MyClass' object has no attribute or method 'name'. Did you forget to initialize an attribute in __init__()?:
File "test-class2.py", line 10
def fn(a: MyClass):
return a.name
~~~~~~ <--- HERE
</pre></div>
</div>
</div>
<div class="section" id="enum-type">
<h3>Enum Type<a class="headerlink" href="#enum-type" title="Permalink to this headline">¶</a></h3>
<p>Like custom classes, semantics of the enum type are user-defined and the entire class definition must be compilable to TorchScript IR and adhere to TorchScript type-checking rules.</p>
<div class="highlight-default notranslate"><div class="highlight"><pre><span></span><span class="n">TSEnumDef</span> <span class="p">:</span><span class="o">:=</span> <span class="s2">"class"</span> <span class="n">Identifier</span> <span class="s2">"(enum.Enum | TSEnumType)"</span> <span class="s2">":"</span>
<span class="p">(</span> <span class="n">MemberIdentifier</span> <span class="s2">"="</span> <span class="n">Value</span> <span class="p">)</span><span class="o">+</span>
<span class="p">(</span> <span class="n">MethodDefinition</span> <span class="p">)</span><span class="o">*</span>
</pre></div>
</div>
<p>Where:</p>
<ul class="simple">
<li><p>Value must be a TorchScript literal of type <code class="docutils literal notranslate"><span class="pre">int</span></code>, <code class="docutils literal notranslate"><span class="pre">float</span></code>, or <code class="docutils literal notranslate"><span class="pre">str</span></code>, and must be of the same TorchScript type.</p></li>
<li><p><code class="docutils literal notranslate"><span class="pre">TSEnumType</span></code> is the name of a TorchScript enumerated type. Similar to Python enum, TorchScript allows restricted <code class="docutils literal notranslate"><span class="pre">Enum</span></code> subclassing, that is, subclassing an enumerated is allowed only if it does not define any members.</p></li>
</ul>
<p><strong>Compared to Python</strong></p>
<ul class="simple">
<li><p>TorchScript supports only <code class="docutils literal notranslate"><span class="pre">enum.Enum</span></code>. It does not support other variations such as <code class="docutils literal notranslate"><span class="pre">enum.IntEnum</span></code>, <code class="docutils literal notranslate"><span class="pre">enum.Flag</span></code>, <code class="docutils literal notranslate"><span class="pre">enum.IntFlag</span></code>, and <code class="docutils literal notranslate"><span class="pre">enum.auto</span></code>.</p></li>
<li><p>Values of TorchScript enum members must be of the same type and can only be <code class="docutils literal notranslate"><span class="pre">int</span></code>, <code class="docutils literal notranslate"><span class="pre">float</span></code>, or <code class="docutils literal notranslate"><span class="pre">str</span></code> types, whereas Python enum members can be of any type.</p></li>
<li><p>Enums containing methods are ignored in TorchScript.</p></li>
</ul>
<p><strong>Example 1</strong></p>
<p>The following example defines the class <code class="docutils literal notranslate"><span class="pre">Color</span></code> as an <code class="docutils literal notranslate"><span class="pre">Enum</span></code> type:</p>
<div class="highlight-default notranslate"><div class="highlight"><pre><span></span><span class="kn">import</span> <span class="nn">torch</span>
<span class="kn">from</span> <span class="nn">enum</span> <span class="kn">import</span> <span class="n">Enum</span>
<span class="k">class</span> <span class="nc">Color</span><span class="p">(</span><span class="n">Enum</span><span class="p">):</span>
<span class="n">RED</span> <span class="o">=</span> <span class="mi">1</span>
<span class="n">GREEN</span> <span class="o">=</span> <span class="mi">2</span>
<span class="k">def</span> <span class="nf">enum_fn</span><span class="p">(</span><span class="n">x</span><span class="p">:</span> <span class="n">Color</span><span class="p">,</span> <span class="n">y</span><span class="p">:</span> <span class="n">Color</span><span class="p">)</span> <span class="o">-></span> <span class="nb">bool</span><span class="p">:</span>
<span class="k">if</span> <span class="n">x</span> <span class="o">==</span> <span class="n">Color</span><span class="o">.</span><span class="n">RED</span><span class="p">:</span>
<span class="k">return</span> <span class="kc">True</span>
<span class="k">return</span> <span class="n">x</span> <span class="o">==</span> <span class="n">y</span>
<span class="n">m</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">jit</span><span class="o">.</span><span class="n">script</span><span class="p">(</span><span class="n">enum_fn</span><span class="p">)</span>
<span class="nb">print</span><span class="p">(</span><span class="s2">"Eager: "</span><span class="p">,</span> <span class="n">enum_fn</span><span class="p">(</span><span class="n">Color</span><span class="o">.</span><span class="n">RED</span><span class="p">,</span> <span class="n">Color</span><span class="o">.</span><span class="n">GREEN</span><span class="p">))</span>
<span class="nb">print</span><span class="p">(</span><span class="s2">"TorchScript: "</span><span class="p">,</span> <span class="n">m</span><span class="p">(</span><span class="n">Color</span><span class="o">.</span><span class="n">RED</span><span class="p">,</span> <span class="n">Color</span><span class="o">.</span><span class="n">GREEN</span><span class="p">))</span>
</pre></div>
</div>
<p><strong>Example 2</strong></p>
<p>The following example shows the case of restricted enum subclassing, where <code class="docutils literal notranslate"><span class="pre">BaseColor</span></code> does not define any member, thus can be subclassed by <code class="docutils literal notranslate"><span class="pre">Color</span></code>:</p>
<div class="highlight-default notranslate"><div class="highlight"><pre><span></span><span class="kn">import</span> <span class="nn">torch</span>
<span class="kn">from</span> <span class="nn">enum</span> <span class="kn">import</span> <span class="n">Enum</span>
<span class="k">class</span> <span class="nc">BaseColor</span><span class="p">(</span><span class="n">Enum</span><span class="p">):</span>
<span class="k">def</span> <span class="nf">foo</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
<span class="k">pass</span>
<span class="k">class</span> <span class="nc">Color</span><span class="p">(</span><span class="n">BaseColor</span><span class="p">):</span>
<span class="n">RED</span> <span class="o">=</span> <span class="mi">1</span>
<span class="n">GREEN</span> <span class="o">=</span> <span class="mi">2</span>
<span class="k">def</span> <span class="nf">enum_fn</span><span class="p">(</span><span class="n">x</span><span class="p">:</span> <span class="n">Color</span><span class="p">,</span> <span class="n">y</span><span class="p">:</span> <span class="n">Color</span><span class="p">)</span> <span class="o">-></span> <span class="nb">bool</span><span class="p">:</span>
<span class="k">if</span> <span class="n">x</span> <span class="o">==</span> <span class="n">Color</span><span class="o">.</span><span class="n">RED</span><span class="p">:</span>
<span class="k">return</span> <span class="kc">True</span>
<span class="k">return</span> <span class="n">x</span> <span class="o">==</span> <span class="n">y</span>
<span class="n">m</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">jit</span><span class="o">.</span><span class="n">script</span><span class="p">(</span><span class="n">enum_fn</span><span class="p">)</span>
<span class="nb">print</span><span class="p">(</span><span class="s2">"TorchScript: "</span><span class="p">,</span> <span class="n">m</span><span class="p">(</span><span class="n">Color</span><span class="o">.</span><span class="n">RED</span><span class="p">,</span> <span class="n">Color</span><span class="o">.</span><span class="n">GREEN</span><span class="p">))</span>
<span class="nb">print</span><span class="p">(</span><span class="s2">"Eager: "</span><span class="p">,</span> <span class="n">enum_fn</span><span class="p">(</span><span class="n">Color</span><span class="o">.</span><span class="n">RED</span><span class="p">,</span> <span class="n">Color</span><span class="o">.</span><span class="n">GREEN</span><span class="p">))</span>
</pre></div>
</div>
</div>
<div class="section" id="torchscript-module-class">
<h3>TorchScript Module Class<a class="headerlink" href="#torchscript-module-class" title="Permalink to this headline">¶</a></h3>
<p><code class="docutils literal notranslate"><span class="pre">TSModuleType</span></code> is a special class type that is inferred from object instances that are created outside TorchScript. <code class="docutils literal notranslate"><span class="pre">TSModuleType</span></code> is named by the Python class of the object instance. The <code class="docutils literal notranslate"><span class="pre">__init__()</span></code> method of the Python class is not considered a TorchScript method, so it does not have to comply with TorchScript’s type-checking rules.</p>
<p>The type schema of a module instance class is constructed directly from an instance object (created outside the scope of TorchScript) rather than inferred from <code class="docutils literal notranslate"><span class="pre">__init__()</span></code> like custom classes. It is possible that two objects of the same instance class type follow two different type schemas.</p>
<p>In this sense, <code class="docutils literal notranslate"><span class="pre">TSModuleType</span></code> is not really a static type. Therefore, for type safety considerations, <code class="docutils literal notranslate"><span class="pre">TSModuleType</span></code> cannot be used in a TorchScript type annotation or be composed with <code class="docutils literal notranslate"><span class="pre">TSType</span></code>.</p>
</div>
<div class="section" id="module-instance-class">
<h3>Module Instance Class<a class="headerlink" href="#module-instance-class" title="Permalink to this headline">¶</a></h3>
<p>TorchScript module type represents the type schema of a user-defined PyTorch module instance. When scripting a PyTorch module, the module object is always created outside TorchScript (i.e., passed in as parameter to <code class="docutils literal notranslate"><span class="pre">forward</span></code>). The Python module class is treated as a module instance class, so the <code class="docutils literal notranslate"><span class="pre">__init__()</span></code> method of the Python module class is not subject to the type-checking rules of TorchScript.</p>
<div class="highlight-default notranslate"><div class="highlight"><pre><span></span><span class="n">TSModuleType</span> <span class="p">:</span><span class="o">:=</span> <span class="s2">"class"</span> <span class="n">Identifier</span> <span class="s2">"(torch.nn.Module)"</span> <span class="s2">":"</span>
<span class="n">ClassBodyDefinition</span>
</pre></div>
</div>
<p>Where:</p>
<ul class="simple">
<li><p><code class="docutils literal notranslate"><span class="pre">forward()</span></code> and other methods decorated with <code class="docutils literal notranslate"><span class="pre">@torch.jit.export</span></code> must be compilable to TorchScript IR and subject to TorchScript’s type-checking rules.</p></li>
</ul>
<p>Unlike custom classes, only the forward method and other methods decorated with <code class="docutils literal notranslate"><span class="pre">@torch.jit.export</span></code> of the module type need to be compilable. Most notably, <code class="docutils literal notranslate"><span class="pre">__init__()</span></code> is not considered a TorchScript method. Consequently, module type constructors cannot be invoked within the scope of TorchScript. Instead, TorchScript module objects are always constructed outside and passed into <code class="docutils literal notranslate"><span class="pre">torch.jit.script(ModuleObj)</span></code>.</p>
<p><strong>Example 1</strong></p>
<p>This example illustrates a few features of module types:</p>
<ul class="simple">
<li><p>The <code class="docutils literal notranslate"><span class="pre">TestModule</span></code> instance is created outside the scope of TorchScript (i.e., before invoking <code class="docutils literal notranslate"><span class="pre">torch.jit.script</span></code>).</p></li>
<li><p><code class="docutils literal notranslate"><span class="pre">__init__()</span></code> is not considered a TorchScript method, therefore, it does not have to be annotated and can contain arbitrary Python code. In addition, the <code class="docutils literal notranslate"><span class="pre">__init__()</span></code> method of an instance class cannot be invoked in TorchScript code. Because <code class="docutils literal notranslate"><span class="pre">TestModule</span></code> instances are instantiated in Python, in this example, <code class="docutils literal notranslate"><span class="pre">TestModule(2.0)</span></code> and <code class="docutils literal notranslate"><span class="pre">TestModule(2)</span></code> create two instances with different types for its data attributes. <code class="docutils literal notranslate"><span class="pre">self.x</span></code> is of type <code class="docutils literal notranslate"><span class="pre">float</span></code> for <code class="docutils literal notranslate"><span class="pre">TestModule(2.0)</span></code>, whereas <code class="docutils literal notranslate"><span class="pre">self.y</span></code> is of type <code class="docutils literal notranslate"><span class="pre">int</span></code> for <code class="docutils literal notranslate"><span class="pre">TestModule(2.0)</span></code>.</p></li>
<li><p>TorchScript automatically compiles other methods (e.g., <code class="docutils literal notranslate"><span class="pre">mul()</span></code>) invoked by methods annotated via <code class="docutils literal notranslate"><span class="pre">@torch.jit.export</span></code> or <code class="docutils literal notranslate"><span class="pre">forward()</span></code> methods.</p></li>
<li><p>Entry-points to a TorchScript program are either <code class="docutils literal notranslate"><span class="pre">forward()</span></code> of a module type, functions annotated as <code class="docutils literal notranslate"><span class="pre">torch.jit.script</span></code>, or methods annotated as <code class="docutils literal notranslate"><span class="pre">torch.jit.export</span></code>.</p></li>
</ul>
<div class="highlight-python3 notranslate"><div class="highlight"><pre><span></span><span class="kn">import</span> <span class="nn">torch</span>
<span class="k">class</span> <span class="nc">TestModule</span><span class="p">(</span><span class="n">torch</span><span class="o">.</span><span class="n">nn</span><span class="o">.</span><span class="n">Module</span><span class="p">):</span>
<span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">v</span><span class="p">):</span>
<span class="nb">super</span><span class="p">()</span><span class="o">.</span><span class="fm">__init__</span><span class="p">()</span>
<span class="bp">self</span><span class="o">.</span><span class="n">x</span> <span class="o">=</span> <span class="n">v</span>
<span class="k">def</span> <span class="nf">forward</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">inc</span><span class="p">:</span> <span class="nb">int</span><span class="p">):</span>
<span class="k">return</span> <span class="bp">self</span><span class="o">.</span><span class="n">x</span> <span class="o">+</span> <span class="n">inc</span>
<span class="n">m</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">jit</span><span class="o">.</span><span class="n">script</span><span class="p">(</span><span class="n">TestModule</span><span class="p">(</span><span class="mi">1</span><span class="p">))</span>
<span class="nb">print</span><span class="p">(</span><span class="sa">f</span><span class="s2">"First instance: </span><span class="si">{</span><span class="n">m</span><span class="p">(</span><span class="mi">3</span><span class="p">)</span><span class="si">}</span><span class="s2">"</span><span class="p">)</span>
<span class="n">m</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">jit</span><span class="o">.</span><span class="n">script</span><span class="p">(</span><span class="n">TestModule</span><span class="p">(</span><span class="n">torch</span><span class="o">.</span><span class="n">ones</span><span class="p">([</span><span class="mi">5</span><span class="p">])))</span>
<span class="nb">print</span><span class="p">(</span><span class="sa">f</span><span class="s2">"Second instance: </span><span class="si">{</span><span class="n">m</span><span class="p">(</span><span class="mi">3</span><span class="p">)</span><span class="si">}</span><span class="s2">"</span><span class="p">)</span>
</pre></div>
</div>
<p>The example above produces the following output:</p>
<div class="highlight-none notranslate"><div class="highlight"><pre><span></span>First instance: 4
Second instance: tensor([4., 4., 4., 4., 4.])
</pre></div>
</div>
<p><strong>Example 2</strong></p>
<p>The following example shows an incorrect usage of module type. Specifically, this example invokes the constructor of <code class="docutils literal notranslate"><span class="pre">TestModule</span></code> inside the scope of TorchScript:</p>
<div class="highlight-python3 notranslate"><div class="highlight"><pre><span></span><span class="kn">import</span> <span class="nn">torch</span>
<span class="k">class</span> <span class="nc">TestModule</span><span class="p">(</span><span class="n">torch</span><span class="o">.</span><span class="n">nn</span><span class="o">.</span><span class="n">Module</span><span class="p">):</span>
<span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">v</span><span class="p">):</span>
<span class="nb">super</span><span class="p">()</span><span class="o">.</span><span class="fm">__init__</span><span class="p">()</span>
<span class="bp">self</span><span class="o">.</span><span class="n">x</span> <span class="o">=</span> <span class="n">v</span>
<span class="k">def</span> <span class="nf">forward</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">x</span><span class="p">:</span> <span class="nb">int</span><span class="p">):</span>
<span class="k">return</span> <span class="bp">self</span><span class="o">.</span><span class="n">x</span> <span class="o">+</span> <span class="n">x</span>
<span class="k">class</span> <span class="nc">MyModel</span><span class="p">:</span>
<span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">v</span><span class="p">:</span> <span class="nb">int</span><span class="p">):</span>
<span class="bp">self</span><span class="o">.</span><span class="n">val</span> <span class="o">=</span> <span class="n">v</span>
<span class="nd">@torch</span><span class="o">.</span><span class="n">jit</span><span class="o">.</span><span class="n">export</span>
<span class="k">def</span> <span class="nf">doSomething</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">val</span><span class="p">:</span> <span class="nb">int</span><span class="p">)</span> <span class="o">-></span> <span class="nb">int</span><span class="p">:</span>
<span class="c1"># error: should not invoke the constructor of module type</span>
<span class="n">myModel</span> <span class="o">=</span> <span class="n">TestModule</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">val</span><span class="p">)</span>
<span class="k">return</span> <span class="n">myModel</span><span class="p">(</span><span class="n">val</span><span class="p">)</span>
<span class="c1"># m = torch.jit.script(MyModel(2)) # Results in below RuntimeError</span>
<span class="c1"># RuntimeError: Could not get name of python class object</span>
</pre></div>
</div>
</div>
</div>
<div class="section" id="type-annotation">
<span id="id2"></span><h2><a class="toc-backref" href="#id9">Type Annotation</a><a class="headerlink" href="#type-annotation" title="Permalink to this headline">¶</a></h2>
<p>Since TorchScript is statically typed, programmers need to annotate types at <em>strategic points</em> of TorchScript code so that every local variable or
instance data attribute has a static type, and every function and method has a statically typed signature.</p>
<div class="section" id="when-to-annotate-types">
<h3>When to Annotate Types<a class="headerlink" href="#when-to-annotate-types" title="Permalink to this headline">¶</a></h3>
<p>In general, type annotations are only needed in places where static types cannot be automatically inferred (e.g., parameters or sometimes return types to
methods or functions). Types of local variables and data attributes are often automatically inferred from their assignment statements. Sometimes an inferred type
may be too restrictive, e.g., <code class="docutils literal notranslate"><span class="pre">x</span></code> being inferred as <code class="docutils literal notranslate"><span class="pre">NoneType</span></code> through assignment <code class="docutils literal notranslate"><span class="pre">x</span> <span class="pre">=</span> <span class="pre">None</span></code>, whereas <code class="docutils literal notranslate"><span class="pre">x</span></code> is actually used as an <code class="docutils literal notranslate"><span class="pre">Optional</span></code>. In such
cases, type annotations may be needed to overwrite auto inference, e.g., <code class="docutils literal notranslate"><span class="pre">x:</span> <span class="pre">Optional[int]</span> <span class="pre">=</span> <span class="pre">None</span></code>. Note that it is always safe to type annotate a local variable
or data attribute even if its type can be automatically inferred. The annotated type must be congruent with TorchScript’s type-checking.</p>
<p>When a parameter, local variable, or data attribute is not type annotated and its type cannot be automatically inferred, TorchScript assumes it to be a
default type of <code class="docutils literal notranslate"><span class="pre">TensorType</span></code>, <code class="docutils literal notranslate"><span class="pre">List[TensorType]</span></code>, or <code class="docutils literal notranslate"><span class="pre">Dict[str,</span> <span class="pre">TensorType]</span></code>.</p>
</div>
<div class="section" id="annotate-function-signature">
<h3>Annotate Function Signature<a class="headerlink" href="#annotate-function-signature" title="Permalink to this headline">¶</a></h3>
<p>Since a parameter may not be automatically inferred from the body of the function (including both functions and methods), they need to be type annotated. Otherwise, they assume the default type <code class="docutils literal notranslate"><span class="pre">TensorType</span></code>.</p>
<p>TorchScript supports two styles for method and function signature type annotation:</p>
<ul class="simple">
<li><p><strong>Python3-style</strong> annotates types directly on the signature. As such, it allows individual parameters to be left unannotated (whose type will be the default type of <code class="docutils literal notranslate"><span class="pre">TensorType</span></code>), or allows the return type to be left unannotated (whose type will be automatically inferred).</p></li>
</ul>
<div class="highlight-default notranslate"><div class="highlight"><pre><span></span><span class="n">Python3Annotation</span> <span class="p">:</span><span class="o">:=</span> <span class="s2">"def"</span> <span class="n">Identifier</span> <span class="p">[</span> <span class="s2">"("</span> <span class="n">ParamAnnot</span><span class="o">*</span> <span class="s2">")"</span> <span class="p">]</span> <span class="p">[</span><span class="n">ReturnAnnot</span><span class="p">]</span> <span class="s2">":"</span>
<span class="n">FuncOrMethodBody</span>
<span class="n">ParamAnnot</span> <span class="p">:</span><span class="o">:=</span> <span class="n">Identifier</span> <span class="p">[</span> <span class="s2">":"</span> <span class="n">TSType</span> <span class="p">]</span> <span class="s2">","</span>
<span class="n">ReturnAnnot</span> <span class="p">:</span><span class="o">:=</span> <span class="s2">"->"</span> <span class="n">TSType</span>
</pre></div>
</div>
<p>Note that when using Python3 style, the type <code class="docutils literal notranslate"><span class="pre">self</span></code> is automatically inferred and should not be annotated.</p>
<ul class="simple">
<li><p><strong>Mypy style</strong> annotates types as a comment right below the function/method declaration. In the Mypy style, since parameter names do not appear in the annotation, all parameters have to be annotated.</p></li>
</ul>
<div class="highlight-default notranslate"><div class="highlight"><pre><span></span><span class="n">MyPyAnnotation</span> <span class="p">:</span><span class="o">:=</span> <span class="s2">"# type:"</span> <span class="s2">"("</span> <span class="n">ParamAnnot</span><span class="o">*</span> <span class="s2">")"</span> <span class="p">[</span> <span class="n">ReturnAnnot</span> <span class="p">]</span>
<span class="n">ParamAnnot</span> <span class="p">:</span><span class="o">:=</span> <span class="n">TSType</span> <span class="s2">","</span>
<span class="n">ReturnAnnot</span> <span class="p">:</span><span class="o">:=</span> <span class="s2">"->"</span> <span class="n">TSType</span>
</pre></div>
</div>
<p><strong>Example 1</strong></p>
<p>In this example:</p>
<ul class="simple">
<li><p><code class="docutils literal notranslate"><span class="pre">a</span></code> is not annotated and assumes the default type of <code class="docutils literal notranslate"><span class="pre">TensorType</span></code>.</p></li>
<li><p><code class="docutils literal notranslate"><span class="pre">b</span></code> is annotated as type <code class="docutils literal notranslate"><span class="pre">int</span></code>.</p></li>
<li><p>The return type is not annotated and is automatically inferred as type <code class="docutils literal notranslate"><span class="pre">TensorType</span></code> (based on the type of the value being returned).</p></li>
</ul>
<div class="highlight-default notranslate"><div class="highlight"><pre><span></span><span class="kn">import</span> <span class="nn">torch</span>
<span class="k">def</span> <span class="nf">f</span><span class="p">(</span><span class="n">a</span><span class="p">,</span> <span class="n">b</span><span class="p">:</span> <span class="nb">int</span><span class="p">):</span>
<span class="k">return</span> <span class="n">a</span><span class="o">+</span><span class="n">b</span>
<span class="n">m</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">jit</span><span class="o">.</span><span class="n">script</span><span class="p">(</span><span class="n">f</span><span class="p">)</span>
<span class="nb">print</span><span class="p">(</span><span class="s2">"TorchScript:"</span><span class="p">,</span> <span class="n">m</span><span class="p">(</span><span class="n">torch</span><span class="o">.</span><span class="n">ones</span><span class="p">([</span><span class="mi">6</span><span class="p">]),</span> <span class="mi">100</span><span class="p">))</span>
</pre></div>
</div>
<p><strong>Example 2</strong></p>
<p>The following example uses Mypy style annotation. Note that parameters or return values must be annotated even if some of
them assume the default type.</p>
<div class="highlight-default notranslate"><div class="highlight"><pre><span></span><span class="kn">import</span> <span class="nn">torch</span>
<span class="k">def</span> <span class="nf">f</span><span class="p">(</span><span class="n">a</span><span class="p">,</span> <span class="n">b</span><span class="p">):</span>
<span class="c1"># type: (torch.Tensor, int) → torch.Tensor</span>
<span class="k">return</span> <span class="n">a</span><span class="o">+</span><span class="n">b</span>