@@ -4547,30 +4547,33 @@ async def test_broadcast_mask_lower_rank(self) -> None:
45474547 )
45484548
45494549
4550- @pytest .mark .parametrize ("dynamic" , [False , True ])
45514550@pytest .mark .parametrize (
4552- "x,dim,index" ,
4551+ "x,dim,index,dynamic " ,
45534552 [
45544553 # 2D float32, select along dim 0
4555- (torch .rand (3 , 4 , dtype = torch .float32 ), 0 , 1 ),
4554+ (torch .rand (3 , 4 , dtype = torch .float32 ), 0 , 1 , False ),
4555+ (torch .rand (3 , 4 , dtype = torch .float32 ), 0 , 1 , True ),
45564556 # 2D float32, select along dim 1
4557- (torch .rand (3 , 4 , dtype = torch .float32 ), 1 , 2 ),
4557+ (torch .rand (3 , 4 , dtype = torch .float32 ), 1 , 2 , False ),
45584558 # 3D float16, select along dim 1
4559- (torch .rand (2 , 3 , 4 , dtype = torch .float16 ), 1 , 2 ),
4559+ (torch .rand (2 , 3 , 4 , dtype = torch .float16 ), 1 , 2 , False ),
45604560 # 3D int32, select along dim 2
4561- (torch .randint (0 , 100 , (2 , 3 , 4 ), dtype = torch .int32 ), 2 , 3 ),
4561+ (torch .randint (0 , 100 , (2 , 3 , 4 ), dtype = torch .int32 ), 2 , 3 , False ),
45624562 # 2D int64, select along dim 0
4563- (torch .randint (- 50 , 50 , (4 , 5 ), dtype = torch .int64 ), 0 , 2 ),
4563+ (torch .randint (- 50 , 50 , (4 , 5 ), dtype = torch .int64 ), 0 , 2 , False ),
45644564 # Negative dimension (dim=-1 is last dim, float32)
4565- (torch .rand (3 , 4 , 5 , dtype = torch .float32 ), - 1 , 2 ),
4565+ (torch .rand (3 , 4 , 5 , dtype = torch .float32 ), - 1 , 2 , False ),
45664566 # Negative index (index from end, int32) — exercises dynamic path when dynamic=True
4567- (torch .randint (0 , 100 , (4 , 5 ), dtype = torch .int32 ), 0 , - 1 ),
4567+ (torch .randint (0 , 100 , (4 , 5 ), dtype = torch .int32 ), 0 , - 1 , False ),
4568+ (torch .randint (0 , 100 , (4 , 5 ), dtype = torch .int32 ), 0 , - 1 , True ),
45684569 # 1D tensor select (float32)
4569- (torch .rand (10 , dtype = torch .float32 ), 0 , 5 ),
4570+ (torch .rand (10 , dtype = torch .float32 ), 0 , 5 , False ),
45704571 # 1D tensor, negative index — exercises 1D dynamic path when dynamic=True
4571- (torch .rand (8 , dtype = torch .float32 ), 0 , - 2 ),
4572+ (torch .rand (8 , dtype = torch .float32 ), 0 , - 2 , False ),
4573+ (torch .rand (8 , dtype = torch .float32 ), 0 , - 2 , True ),
45724574 # 3D float32, negative index on last dim — exercises dynamic path when dynamic=True
4573- (torch .rand (2 , 3 , 5 , dtype = torch .float32 ), 2 , - 1 ),
4575+ (torch .rand (2 , 3 , 5 , dtype = torch .float32 ), 2 , - 1 , False ),
4576+ (torch .rand (2 , 3 , 5 , dtype = torch .float32 ), 2 , - 1 , True ),
45744577 ],
45754578)
45764579async def test_select_int (x : Tensor , dim : int , index : int , dynamic : bool ) -> None :
0 commit comments