refactor: simplify forward() permutation logic for compile-friendly execution

What's changed
- Replace tensor-based perm construction with list-based version
- Add explicit inverse permutation for correctness
- Remove TorchScript-specific branches

Why
- Compile-friendly: torch.compile/AOTAutograd prefer static Python control flow and index lists over device tensor construction inside forward. Replacing torch.tensor([...]), torch.arange(...), and torch.cat(...) with plain Python lists reduces graph breaks and guard complexity, improving compilation stability and cache reuse.

Author: studyingeugene

compressai/entropy_models/entropy_models.py

@@ -474,28 +474,18 @@ def forward(
         if training is None:
             training = self.training
 
-        if not torch.jit.is_scripting():
-            # x from B x C x ... to C x B x ...
-            perm = torch.cat(
-                (
-                    torch.tensor([1, 0], dtype=torch.long, device=x.device),
-                    torch.arange(2, x.ndim, dtype=torch.long, device=x.device),
-                )
-            )
-            inv_perm = perm
-        else:
-            raise NotImplementedError()
-            # TorchScript in 2D for static inference
-            # Convert to (channels, ... , batch) format
-            # perm = (1, 2, 3, 0)
-            # inv_perm = (3, 0, 1, 2)
+        D = x.dim()  
+        # B C ...  ->  C B ...
+        perm = [1, 0] + list(range(2, D))
+        inv_perm = [0] * D
+        for i, p in enumerate(perm):
+            inv_perm[p] = i
 
         x = x.permute(*perm).contiguous()
         shape = x.size()
         values = x.reshape(x.size(0), 1, -1)
 
         # Add noise or quantize
-
         outputs = self.quantize(
             values, "noise" if training else "dequantize", self._get_medians()
         )
@@ -510,11 +500,8 @@ def forward(
             # likelihood = torch.zeros_like(outputs)
 
         # Convert back to input tensor shape
-        outputs = outputs.reshape(shape)
-        outputs = outputs.permute(*inv_perm).contiguous()
-
-        likelihood = likelihood.reshape(shape)
-        likelihood = likelihood.permute(*inv_perm).contiguous()
+        outputs = outputs.reshape(shape).permute(*inv_perm).contiguous()
+        likelihood = likelihood.reshape(shape).permute(*inv_perm).contiguous()
 
         return outputs, likelihood