import os os.environ["TRITON_BACKENDS_IN_TREE"] = "1"

import torch import torch.nn as nn

class Model(nn.Module): def init(self): super().init() self.embedding = nn.Embedding(num_embeddings=100, embedding_dim=32) self.linear = nn.Linear(32, 16)

def forward(self, x):
    embedded = self.embedding(x)
    embedded = embedded.float()
    base_tensor = self.linear(embedded)

    # This makes dim=1 have size 1
    base_tensor = base_tensor.unsqueeze(1)  # shape: [batch, 1, seq_len, 16]

    # select(dim=1, index=0) is valid
    # select(dim=1, index=1) is OUT OF BOUNDS — size is 1!
    # select(dim=1, index=2) is OUT OF BOUNDS — size is 1!
    slice0 = torch.select(base_tensor, dim=1, index=0)
    slice1 = torch.select(base_tensor, dim=1, index=1)  # <- should FAIL
    slice2 = torch.select(base_tensor, dim=1, index=2)  # <- should FAIL

    unsqueezed0 = torch.unsqueeze(slice0, dim=1)
    unsqueezed1 = torch.unsqueeze(slice1, dim=1)
    unsqueezed2 = torch.unsqueeze(slice2, dim=1)

    result = torch.cat([unsqueezed0, unsqueezed1, unsqueezed2], dim=1)
    return result

device = "cuda" model = Model().to(device).eval() x = torch.randint(0, 100, (4, 10), dtype=torch.long, device=device)

Eager: should raise IndexError

try: with torch.no_grad(): eager_out = model(x) print(f"eager: OK shape={eager_out.shape}") except (IndexError, RuntimeError) as e: print(f"eager: ERROR — {e}")

Expected: eager: ERROR — select(): index 1 out of range...

Compiled: should also raise, but SILENTLY SUCCEEDS

torch._dynamo.reset() compiled = torch.compile(model, backend="inductor") try: with torch.no_grad(): comp_out = compiled(x) print(f"compile: OK shape={comp_out.shape}") except Exception as e: print(f"compile: ERROR — {e}")

Actual: compile: OK — should have raised the same error!

TL;DR

The issue can be mitigated by avoiding the use of torch.compile with the inductor backend for models that perform out-of-bounds torch.select operations, or by manually adding bounds checks before the select operations.

Guidance

• Identify and review all torch.select operations in the model to ensure that the indices are within the valid range of the tensor's dimensions.
• Consider adding manual bounds checks before the select operations to raise an error if the index is out of range.
• If possible, avoid using the inductor backend with torch.compile for models that perform out-of-bounds torch.select operations, and instead use eager mode or a different backend.
• Verify that the model behaves correctly in eager mode before compiling it with the inductor backend.

Example

# Manual bounds check before torch.select
if index >= base_tensor.size(1):
    raise IndexError(f"select(): index {index} out of range for tensor of size {base_tensor.size()} at dimension 1")
slice1 = torch.select(base_tensor, dim=1, index=1)

Notes

• The root cause of the issue is likely due to the merge_unbind_stack_aten optimization pass in the inductor backend, which bypasses the runtime bounds check for torch.select operations.
• The issue may be specific to the inductor backend and may not occur with other backends or in eager mode.

Recommendation

Apply a workaround by manually adding bounds checks before the select operations, as the inductor backend may not be correctly handling out-of-bounds indices. This will ensure that the model raises an error if an out-of-bounds index is used, rather than silently producing an incorrect output.