Uploaded image for project: 'AI Platform Core Components'
  1. AI Platform Core Components
  2. AIPCC-8537

Inductor fails to compile models using `torch.cat` with `axis` keyword argument

    • Icon: Bug Bug
    • Resolution: Done
    • Icon: Undefined Undefined
    • None
    • None
    • PyTorch
    • 2
    • PyTorch Sprint 23, PyTorch Sprint 24, PyTorch Sprint 25, PyTorch Sprint 26

          1. 🐛 Describe the bug

      When compiling a model with `torch.compile(backend="inductor")`, if the model's `forward` method uses `torch.cat` with the keyword argument `axis` (e.g., `torch.cat([x, t], axis=1)`), compilation fails with a `TypeError`. The error indicates that an internal Inductor function does not recognize the `axis` parameter name.

      The same model compiles and runs successfully with `backend="eager"` and `backend="aot_eager"`. The PyTorch API documents `axis` as a valid alias for `dim` in `torch.cat`, so this is an inconsistency within the Inductor compiler's optimization passes.

          1. To reproduce
            ```python
            import torch.nn as nn
            import torch
            from torch import nn

      class TestModel(nn.Module):

      def _init_(self):
      super(TestModel, self)._init_()
      self.stack = nn.Sequential(
      nn.Linear(256, 128),
      nn.Linear(128, 64),
      nn.Linear(64, 10)
      )

      def forward(self, x, t):
      X = torch.cat([x, t], axis=1)
      X = torch.relu(X)
      return X

      input1 = torch.randn(16, 200)
      input2 = torch.randn(16, 56)
      model = TestModel()

      def run_test(model, input, backend):
      try:
      model = torch.compile(model, backend=backend)
      output = model(*input)
      print(f"succeed on

      {backend}")
      except Exception as e:
      print(f"failed on {backend}

      :", str(e))

      run_test(model, [input1, input2], "eager")
      run_test(model, [input1, input2], "aot_eager")
      run_test(model, [input1, input2], "inductor")
      ```

          1. Error logs

      ```
      succeed on eager
      succeed on aot_eager
      failed on inductor: backend='inductor' raised:
      TypeError: sink_cat_after_pointwise.<locals>.cat_args() got an unexpected keyword argument 'axis'
      ```

          1. Versions

      Collecting environment information...
      PyTorch version: 2.11.0.dev20260107+cpu
      Is debug build: False
      CUDA used to build PyTorch: None
      ROCM used to build PyTorch: N/A

      OS: Microsoft Windows 11
      GCC version: Could not collect
      Clang version: Could not collect
      CMake version: version 4.0.2
      Libc version: N/A

      Python version: 3.10.10 (tags/v3.10.10:aad5f6a, Feb 7 2023, 17:20:36) [MSC v.1929 64 bit (AMD64)] (64-bit runtime)
      Python platform: Windows-10-10.0.26200-SP0
      Is CUDA available: False
      CUDA runtime version: No CUDA
      CUDA_MODULE_LOADING set to: N/A
      GPU models and configuration: No CUDA
      Nvidia driver version: No CUDA
      cuDNN version: No CUDA
      HIP runtime version: N/A
      MIOpen runtime version: N/A
      Is XNNPACK available: True

      cc @chauhang @penguinwu @voznesenskym @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @wenzhe-nrv @jiayisunx @ipiszy @kadeng @muchulee8 @amjames @aakhundov @coconutruben @jataylo

              rh-ee-hgudmund Hinrik Gudmundsson
              rh-ee-hgudmund Hinrik Gudmundsson
              PyTorch Compile
              Votes:
              0 Vote for this issue
              Watchers:
              1 Start watching this issue

                Created:
                Updated:
                Resolved: