diff --git a/modules/sdnq/packed_int.py b/modules/sdnq/packed_int.py index c93cbc7e9..31e8d9214 100644 --- a/modules/sdnq/packed_int.py +++ b/modules/sdnq/packed_int.py @@ -27,28 +27,45 @@ def unpack_int_asymetric(packed_tensor: torch.ByteTensor, shape: torch.Size, wei def pack_uint7(tensor: torch.ByteTensor) -> torch.ByteTensor: packed_tensor = tensor.contiguous().view(-1, 8) - packed_tensor = torch.stack( - ( - torch.bitwise_or(packed_tensor[:, 0], torch.bitwise_and(torch.bitwise_left_shift(packed_tensor[:, 7], 1), 128)), - torch.bitwise_or(packed_tensor[:, 1], torch.bitwise_and(torch.bitwise_left_shift(packed_tensor[:, 7], 2), 128)), - torch.bitwise_or(packed_tensor[:, 2], torch.bitwise_and(torch.bitwise_left_shift(packed_tensor[:, 7], 3), 128)), - torch.bitwise_or(packed_tensor[:, 3], torch.bitwise_and(torch.bitwise_left_shift(packed_tensor[:, 7], 4), 128)), - torch.bitwise_or(packed_tensor[:, 4], torch.bitwise_and(torch.bitwise_left_shift(packed_tensor[:, 7], 5), 128)), - torch.bitwise_or(packed_tensor[:, 5], torch.bitwise_and(torch.bitwise_left_shift(packed_tensor[:, 7], 6), 128)), - torch.bitwise_or(packed_tensor[:, 6], torch.bitwise_and(torch.bitwise_left_shift(packed_tensor[:, 7], 7), 128)), + packed_tensor = torch.bitwise_or( + packed_tensor[:, :7], + torch.bitwise_and( + torch.stack( + ( + torch.bitwise_left_shift(packed_tensor[:, 7], 1), + torch.bitwise_left_shift(packed_tensor[:, 7], 2), + torch.bitwise_left_shift(packed_tensor[:, 7], 3), + torch.bitwise_left_shift(packed_tensor[:, 7], 4), + torch.bitwise_left_shift(packed_tensor[:, 7], 5), + torch.bitwise_left_shift(packed_tensor[:, 7], 6), + torch.bitwise_left_shift(packed_tensor[:, 7], 7), + ), + dim=-1 + ), + 128 ), - dim=-1 ) return packed_tensor def pack_uint6(tensor: torch.ByteTensor) -> torch.ByteTensor: packed_tensor = tensor.contiguous().view(-1, 4) - packed_tensor = torch.stack( + packed_tensor = torch.cat( ( - torch.bitwise_or(packed_tensor[:, 0], torch.bitwise_and(torch.bitwise_left_shift(packed_tensor[:, 3], 2), 192)), - torch.bitwise_or(packed_tensor[:, 1], torch.bitwise_and(torch.bitwise_left_shift(packed_tensor[:, 3], 4), 192)), - torch.bitwise_or(packed_tensor[:, 2], torch.bitwise_left_shift(packed_tensor[:, 3], 6)), + torch.bitwise_or( + packed_tensor[:, :2], + torch.bitwise_and( + torch.stack( + ( + torch.bitwise_left_shift(packed_tensor[:, 3], 2), + torch.bitwise_left_shift(packed_tensor[:, 3], 4), + ), + dim=-1 + ), + 192 + ) + ), + torch.bitwise_or(packed_tensor[:, 2], torch.bitwise_left_shift(packed_tensor[:, 3], 6)).unsqueeze(-1), ), dim=-1 ) @@ -57,25 +74,23 @@ def pack_uint6(tensor: torch.ByteTensor) -> torch.ByteTensor: def pack_uint5(tensor: torch.ByteTensor) -> torch.ByteTensor: packed_tensor = tensor.contiguous().view(-1, 8) - packed_tensor = torch.stack( + packed_tensor = torch.cat( ( - torch.bitwise_or(packed_tensor[:, 0], torch.bitwise_left_shift(packed_tensor[:, 5], 5)), - torch.bitwise_or(packed_tensor[:, 1], torch.bitwise_left_shift(packed_tensor[:, 6], 5)), - torch.bitwise_or(packed_tensor[:, 2], torch.bitwise_left_shift(packed_tensor[:, 7], 5)), + torch.bitwise_or(packed_tensor[:, :3], torch.bitwise_left_shift(packed_tensor[:, 5:8], 5)), torch.bitwise_or( packed_tensor[:, 3], torch.bitwise_or( torch.bitwise_and(torch.bitwise_left_shift(packed_tensor[:, 5], 2), 96), torch.bitwise_and(torch.bitwise_left_shift(packed_tensor[:, 7], 3), 128), ), - ), + ).unsqueeze(-1), torch.bitwise_or( packed_tensor[:, 4], torch.bitwise_or( torch.bitwise_and(torch.bitwise_left_shift(packed_tensor[:, 6], 2), 96), torch.bitwise_and(torch.bitwise_left_shift(packed_tensor[:, 7], 4), 128), ), - ), + ).unsqueeze(-1), ), dim=-1 ) @@ -90,25 +105,18 @@ def pack_uint4(tensor: torch.ByteTensor) -> torch.ByteTensor: def pack_uint3(tensor: torch.ByteTensor) -> torch.ByteTensor: packed_tensor = tensor.contiguous().view(-1, 8) - packed_tensor = torch.stack( - ( - torch.bitwise_or( - torch.bitwise_or(packed_tensor[:, 0], torch.bitwise_left_shift(packed_tensor[:, 1], 3)), - torch.bitwise_left_shift(packed_tensor[:, 6], 6), - ), - torch.bitwise_or( - torch.bitwise_or(packed_tensor[:, 2], torch.bitwise_left_shift(packed_tensor[:, 3], 3)), - torch.bitwise_left_shift(packed_tensor[:, 7], 6), - ), - torch.bitwise_or( - torch.bitwise_or(packed_tensor[:, 4], torch.bitwise_left_shift(packed_tensor[:, 5], 3)), + packed_tensor = torch.bitwise_or( + torch.bitwise_or(packed_tensor[:, :3], torch.bitwise_left_shift(packed_tensor[:, 3:6], 3)), + torch.cat( + ( + torch.bitwise_left_shift(packed_tensor[:, 6:8], 6), torch.bitwise_or( torch.bitwise_and(torch.bitwise_left_shift(packed_tensor[:, 6], 4), 64), torch.bitwise_and(torch.bitwise_left_shift(packed_tensor[:, 7], 5), 128), - ) + ).unsqueeze(-1), ), - ), - dim=-1 + dim=-1 + ) ) return packed_tensor @@ -123,7 +131,7 @@ def pack_uint2(tensor: torch.ByteTensor) -> torch.ByteTensor: def pack_uint1(tensor: torch.Tensor) -> torch.Tensor: - packed_tensor = tensor.contiguous().reshape(-1, 8) + packed_tensor = tensor.contiguous().view(-1, 8) packed_tensor = torch.bitwise_or( torch.bitwise_or( torch.bitwise_or(packed_tensor[:, 0], torch.bitwise_left_shift(packed_tensor[:, 1], 1)), @@ -138,15 +146,9 @@ def pack_uint1(tensor: torch.Tensor) -> torch.Tensor: def unpack_uint7(packed_tensor: torch.ByteTensor, shape: torch.Size) -> torch.ByteTensor: - result = torch.stack( + result = torch.cat( ( - torch.bitwise_and(packed_tensor[:, 0], 127), - torch.bitwise_and(packed_tensor[:, 1], 127), - torch.bitwise_and(packed_tensor[:, 2], 127), - torch.bitwise_and(packed_tensor[:, 3], 127), - torch.bitwise_and(packed_tensor[:, 4], 127), - torch.bitwise_and(packed_tensor[:, 5], 127), - torch.bitwise_and(packed_tensor[:, 6], 127), + torch.bitwise_and(packed_tensor[:, :7], 127), torch.bitwise_or( torch.bitwise_or( torch.bitwise_or( @@ -165,7 +167,7 @@ def unpack_uint7(packed_tensor: torch.ByteTensor, shape: torch.Size) -> torch.By ), torch.bitwise_right_shift(packed_tensor[:, 6], 7), ), - ) + ).unsqueeze(-1) ), dim=-1 ).view(shape) @@ -173,18 +175,16 @@ def unpack_uint7(packed_tensor: torch.ByteTensor, shape: torch.Size) -> torch.By def unpack_uint6(packed_tensor: torch.ByteTensor, shape: torch.Size) -> torch.ByteTensor: - result = torch.stack( + result = torch.cat( ( - torch.bitwise_and(packed_tensor[:, 0], 63), - torch.bitwise_and(packed_tensor[:, 1], 63), - torch.bitwise_and(packed_tensor[:, 2], 63), + torch.bitwise_and(packed_tensor[:, 0:3], 63), torch.bitwise_or( torch.bitwise_or( torch.bitwise_and(torch.bitwise_right_shift(packed_tensor[:, 0], 2), 48), torch.bitwise_and(torch.bitwise_right_shift(packed_tensor[:, 1], 4), 12), ), torch.bitwise_right_shift(packed_tensor[:, 2], 6) - ) + ).unsqueeze(-1) ), dim=-1 ).view(shape) @@ -192,28 +192,21 @@ def unpack_uint6(packed_tensor: torch.ByteTensor, shape: torch.Size) -> torch.By def unpack_uint5(packed_tensor: torch.ByteTensor, shape: torch.Size) -> torch.ByteTensor: - result = torch.stack( + result_bitwise_right_shift = torch.bitwise_right_shift(packed_tensor[:, :3], 5) + result = torch.cat( ( - torch.bitwise_and(packed_tensor[:, 0], 31), - torch.bitwise_and(packed_tensor[:, 1], 31), - torch.bitwise_and(packed_tensor[:, 2], 31), - torch.bitwise_and(packed_tensor[:, 3], 31), - torch.bitwise_and(packed_tensor[:, 4], 31), + torch.bitwise_and(packed_tensor[:, :5], 31), torch.bitwise_or( - torch.bitwise_right_shift(packed_tensor[:, 0], 5), - torch.bitwise_and(torch.bitwise_right_shift(packed_tensor[:, 3], 2), 24), + result_bitwise_right_shift[:, :2], + torch.bitwise_and(torch.bitwise_right_shift(packed_tensor[:, 3:5], 2), 24), ), torch.bitwise_or( - torch.bitwise_right_shift(packed_tensor[:, 1], 5), - torch.bitwise_and(torch.bitwise_right_shift(packed_tensor[:, 4], 2), 24), - ), - torch.bitwise_or( - torch.bitwise_right_shift(packed_tensor[:, 2], 5), + result_bitwise_right_shift[:, 2], torch.bitwise_or( torch.bitwise_and(torch.bitwise_right_shift(packed_tensor[:, 3], 3), 16), torch.bitwise_and(torch.bitwise_right_shift(packed_tensor[:, 4], 4), 8), ), - ), + ).unsqueeze(-1), ), dim=-1 ).view(shape) @@ -226,21 +219,30 @@ def unpack_uint4(packed_tensor: torch.ByteTensor, shape: torch.Size) -> torch.By def unpack_uint3(packed_tensor: torch.ByteTensor, shape: torch.Size) -> torch.ByteTensor: - result = torch.stack( + result = torch.cat( ( - torch.bitwise_and(packed_tensor[:, 0], 7), - torch.bitwise_and(torch.bitwise_right_shift(packed_tensor[:, 0], 3), 7), - torch.bitwise_and(packed_tensor[:, 1], 7), - torch.bitwise_and(torch.bitwise_right_shift(packed_tensor[:, 1], 3), 7), - torch.bitwise_and(packed_tensor[:, 2], 7), - torch.bitwise_and(torch.bitwise_right_shift(packed_tensor[:, 2], 3), 7), - torch.bitwise_or( - torch.bitwise_right_shift(packed_tensor[:, 0], 6), - torch.bitwise_and(torch.bitwise_right_shift(packed_tensor[:, 2], 4), 4), + torch.bitwise_and( + torch.cat( + ( + packed_tensor[:, :3], + torch.bitwise_right_shift(packed_tensor[:, :3], 3) + ), + dim=-1 + ), + 7 ), torch.bitwise_or( - torch.bitwise_right_shift(packed_tensor[:, 1], 6), - torch.bitwise_and(torch.bitwise_right_shift(packed_tensor[:, 2], 5), 4), + torch.bitwise_right_shift(packed_tensor[:, :2], 6), + torch.bitwise_and( + torch.stack( + ( + torch.bitwise_right_shift(packed_tensor[:, 2], 4), + torch.bitwise_right_shift(packed_tensor[:, 2], 5), + ), + dim=-1 + ), + 4 + ), ), ), dim=-1 @@ -249,12 +251,20 @@ def unpack_uint3(packed_tensor: torch.ByteTensor, shape: torch.Size) -> torch.By def unpack_uint2(packed_tensor: torch.ByteTensor, shape: torch.Size) -> torch.ByteTensor: - result = torch.stack( + result = torch.cat( ( - torch.bitwise_and(packed_tensor, 3), - torch.bitwise_and(torch.bitwise_right_shift(packed_tensor, 2), 3), - torch.bitwise_and(torch.bitwise_right_shift(packed_tensor, 4), 3), - torch.bitwise_right_shift(packed_tensor, 6), + torch.bitwise_and( + torch.stack( + ( + packed_tensor, + torch.bitwise_right_shift(packed_tensor, 2), + torch.bitwise_right_shift(packed_tensor, 4) + ), + dim=-1 + ), + 3 + ), + torch.bitwise_right_shift(packed_tensor, 6).unsqueeze(-1), ), dim=-1 ).view(shape) @@ -262,19 +272,27 @@ def unpack_uint2(packed_tensor: torch.ByteTensor, shape: torch.Size) -> torch.By def unpack_uint1(packed_tensor: torch.Tensor, shape: torch.Size) -> torch.Tensor: - result = torch.stack( + result = torch.cat( ( - torch.bitwise_and(packed_tensor, 1), - torch.bitwise_and(torch.bitwise_right_shift(packed_tensor, 1), 1), - torch.bitwise_and(torch.bitwise_right_shift(packed_tensor, 2), 1), - torch.bitwise_and(torch.bitwise_right_shift(packed_tensor, 3), 1), - torch.bitwise_and(torch.bitwise_right_shift(packed_tensor, 4), 1), - torch.bitwise_and(torch.bitwise_right_shift(packed_tensor, 5), 1), - torch.bitwise_and(torch.bitwise_right_shift(packed_tensor, 6), 1), - torch.bitwise_right_shift(packed_tensor, 7), + torch.bitwise_and( + torch.stack( + ( + packed_tensor, + torch.bitwise_right_shift(packed_tensor, 1), + torch.bitwise_right_shift(packed_tensor, 2), + torch.bitwise_right_shift(packed_tensor, 3), + torch.bitwise_right_shift(packed_tensor, 4), + torch.bitwise_right_shift(packed_tensor, 5), + torch.bitwise_right_shift(packed_tensor, 6), + ), + dim=-1 + ), + 1 + ), + torch.bitwise_right_shift(packed_tensor, 7).unsqueeze(-1), ), dim=-1 - ).reshape(shape) + ).view(shape) return result