@@ -520,6 +520,112 @@ def forward(self, x):
520520 )
521521 _testing .assert_onnx_program (onnx_program )
522522
523+ def test_unbind_dim0 (self ):
524+ """Test unbind along dimension 0"""
525+
526+ class UnbindModel (torch .nn .Module ):
527+ def forward (self , x ):
528+ tensors = torch .unbind (x , dim = 0 )
529+ return sum (tensors )
530+
531+ model = UnbindModel ()
532+ x = torch .randn (3 , 4 , 5 )
533+ onnx_program = torch .onnx .export (model , (x ,), dynamo = True , verbose = False )
534+ _testing .assert_onnx_program (onnx_program )
535+
536+ def test_unbind_dim1 (self ):
537+ """Test unbind along dimension 1"""
538+
539+ class UnbindModel (torch .nn .Module ):
540+ def forward (self , x ):
541+ tensors = torch .unbind (x , dim = 1 )
542+ return sum (tensors )
543+
544+ model = UnbindModel ()
545+ x = torch .randn (2 , 3 , 4 )
546+ onnx_program = torch .onnx .export (model , (x ,), dynamo = True , verbose = False )
547+ _testing .assert_onnx_program (onnx_program )
548+
549+ def test_unbind_negative_dim (self ):
550+ """Test unbind with negative dimension"""
551+
552+ class UnbindModel (torch .nn .Module ):
553+ def forward (self , x ):
554+ tensors = torch .unbind (x , dim = - 1 )
555+ return sum (tensors )
556+
557+ model = UnbindModel ()
558+ x = torch .randn (2 , 3 , 4 )
559+ onnx_program = torch .onnx .export (model , (x ,), dynamo = True , verbose = False )
560+ _testing .assert_onnx_program (onnx_program )
561+
562+ def test_unbind_size_one (self ):
563+ """Test unbind with dimension of size 1"""
564+
565+ class UnbindModel (torch .nn .Module ):
566+ def forward (self , x ):
567+ tensors = torch .unbind (x , dim = 0 )
568+ return tensors [0 ]
569+
570+ model = UnbindModel ()
571+ x = torch .randn (1 , 4 , 5 )
572+ onnx_program = torch .onnx .export (model , (x ,), dynamo = True , verbose = False )
573+ _testing .assert_onnx_program (onnx_program )
574+
575+ def test_unbind_with_lstm (self ):
576+ """Test unbind in LSTM context"""
577+
578+ class LSTMDecoder (torch .nn .Module ):
579+ def __init__ (self ):
580+ super ().__init__ ()
581+ self .embedding = torch .nn .Embedding (100 , 64 )
582+ self .lstm = torch .nn .LSTM (64 , 64 , 2 , batch_first = True ) # 2 layers
583+ self .fc = torch .nn .Linear (64 , 100 )
584+
585+ def forward (self , tokens , h , c ):
586+ embedded = self .embedding (tokens ).unsqueeze (0 )
587+ output , (h_out , c_out ) = self .lstm (embedded , (h , c ))
588+ logits = self .fc (output .squeeze (0 ).squeeze (0 ))
589+ return logits , h_out , c_out
590+
591+ model = LSTMDecoder ()
592+ model .eval ()
593+ tokens = torch .tensor ([1 ])
594+ h = torch .randn (2 , 1 , 64 ) # 2 layers
595+ c = torch .randn (2 , 1 , 64 ) # 2 layers
596+ onnx_program = torch .onnx .export (model , (tokens , h , c ), dynamo = True , verbose = False )
597+ _testing .assert_onnx_program (onnx_program )
598+
599+ def test_unbind_dynamic_dim0 (self ):
600+ """Test unbind with dynamic dimension 0 - triggers SplitToSequence"""
601+
602+ class UnbindModel (torch .nn .Module ):
603+ def forward (self , x ):
604+ tensors = torch .unbind (x , dim = 0 )
605+ return sum (tensors )
606+
607+ model = UnbindModel ()
608+ x = torch .randn (3 , 4 , 5 )
609+ onnx_program = torch .onnx .export (
610+ model , (x ,), dynamo = True , verbose = False , dynamic_shapes = ({0 : "batch_size" },)
611+ )
612+ _testing .assert_onnx_program (onnx_program )
613+
614+ def test_unbind_dynamic_dim1 (self ):
615+ """Test unbind with dynamic dimension 1 - triggers SplitToSequence"""
616+
617+ class UnbindModel (torch .nn .Module ):
618+ def forward (self , x ):
619+ tensors = torch .unbind (x , dim = 1 )
620+ return sum (tensors )
621+
622+ model = UnbindModel ()
623+ x = torch .randn (2 , 3 , 4 )
624+ onnx_program = torch .onnx .export (
625+ model , (x ,), dynamo = True , verbose = False , dynamic_shapes = ({1 : "seq_len" },)
626+ )
627+ _testing .assert_onnx_program (onnx_program )
628+
523629
524630if __name__ == "__main__" :
525631 unittest .main ()
0 commit comments