@@ -4549,30 +4549,33 @@ async def test_broadcast_mask_lower_rank(self) -> None:
45494549 )
45504550
45514551
4552- @pytest .mark .parametrize ("dynamic" , [False , True ])
45534552@pytest .mark .parametrize (
4554- "x,dim,index" ,
4553+ "x,dim,index,dynamic " ,
45554554 [
45564555 # 2D float32, select along dim 0
4557- (torch .rand (3 , 4 , dtype = torch .float32 ), 0 , 1 ),
4556+ (torch .rand (3 , 4 , dtype = torch .float32 ), 0 , 1 , False ),
4557+ (torch .rand (3 , 4 , dtype = torch .float32 ), 0 , 1 , True ),
45584558 # 2D float32, select along dim 1
4559- (torch .rand (3 , 4 , dtype = torch .float32 ), 1 , 2 ),
4559+ (torch .rand (3 , 4 , dtype = torch .float32 ), 1 , 2 , False ),
45604560 # 3D float16, select along dim 1
4561- (torch .rand (2 , 3 , 4 , dtype = torch .float16 ), 1 , 2 ),
4561+ (torch .rand (2 , 3 , 4 , dtype = torch .float16 ), 1 , 2 , False ),
45624562 # 3D int32, select along dim 2
4563- (torch .randint (0 , 100 , (2 , 3 , 4 ), dtype = torch .int32 ), 2 , 3 ),
4563+ (torch .randint (0 , 100 , (2 , 3 , 4 ), dtype = torch .int32 ), 2 , 3 , False ),
45644564 # 2D int64, select along dim 0
4565- (torch .randint (- 50 , 50 , (4 , 5 ), dtype = torch .int64 ), 0 , 2 ),
4565+ (torch .randint (- 50 , 50 , (4 , 5 ), dtype = torch .int64 ), 0 , 2 , False ),
45664566 # Negative dimension (dim=-1 is last dim, float32)
4567- (torch .rand (3 , 4 , 5 , dtype = torch .float32 ), - 1 , 2 ),
4567+ (torch .rand (3 , 4 , 5 , dtype = torch .float32 ), - 1 , 2 , False ),
45684568 # Negative index (index from end, int32) — exercises dynamic path when dynamic=True
4569- (torch .randint (0 , 100 , (4 , 5 ), dtype = torch .int32 ), 0 , - 1 ),
4569+ (torch .randint (0 , 100 , (4 , 5 ), dtype = torch .int32 ), 0 , - 1 , False ),
4570+ (torch .randint (0 , 100 , (4 , 5 ), dtype = torch .int32 ), 0 , - 1 , True ),
45704571 # 1D tensor select (float32)
4571- (torch .rand (10 , dtype = torch .float32 ), 0 , 5 ),
4572+ (torch .rand (10 , dtype = torch .float32 ), 0 , 5 , False ),
45724573 # 1D tensor, negative index — exercises 1D dynamic path when dynamic=True
4573- (torch .rand (8 , dtype = torch .float32 ), 0 , - 2 ),
4574+ (torch .rand (8 , dtype = torch .float32 ), 0 , - 2 , False ),
4575+ (torch .rand (8 , dtype = torch .float32 ), 0 , - 2 , True ),
45744576 # 3D float32, negative index on last dim — exercises dynamic path when dynamic=True
4575- (torch .rand (2 , 3 , 5 , dtype = torch .float32 ), 2 , - 1 ),
4577+ (torch .rand (2 , 3 , 5 , dtype = torch .float32 ), 2 , - 1 , False ),
4578+ (torch .rand (2 , 3 , 5 , dtype = torch .float32 ), 2 , - 1 , True ),
45764579 ],
45774580)
45784581async def test_select_int (x : Tensor , dim : int , index : int , dynamic : bool ) -> None :
0 commit comments