Support for squeeze and select operators in ExecuTorch

[vikasbalaga]

🚀 The feature, motivation and pitch

Hi,

I have been working on a torch model that is based on "state space representation" and trying to export it to ExecuTorch and perform inference.

But during export(), I have faced some issues, due to the presence of a python "for loop", in my forward(), which undergoes loop unrolling and resulted in abnormal model size.

For that, I have used the PyTorch "while_loop" operator and with the help of @angelayi , able to perform "Edge IR" conversion of the model.

Post Edge IR conversion, I am facing the following issue :

WARNING:root:Schema: aten::select.t(t[](a) list, int idx) -> t(*) cannot be converted to torch.FunctionSchema
WARNING:root:Schema: aten::select.t(t[](a) list, int idx) -> t(*) cannot be converted to torch.FunctionSchema
RuntimeError: Missing out variants: {'aten::squeeze', 'aten::select'}

So, I am not sure if these operators are supported in ExecuTorch or not and I have tried many workarounds but all of them boil down to same issue.

So, can someone please help me in identifying the issues here i.e., am I missing something or is the support for {'aten::squeeze', 'aten::select'} not present in ExecuTorch.

If the support is not present, can someone provide help to resolve my issue.

Please find below the sample code which will reproduce the issue.

import torch
from torch import nn
from torch.export import export
from executorch.exir import to_edge_transform_and_lower

class ExportableLoop(nn.Module):
    def __init__(self, hidden_size, out_channels):
        super().__init__()
        self.hidden_size = hidden_size
        self.B = nn.Parameter(torch.randn(hidden_size, 1))  # (H, in_channels)
        self.C = nn.Parameter(torch.randn(out_channels, hidden_size))  # (C_out, H)
        A = torch.randn(2, hidden_size)
        self.A_real = nn.Parameter(A[0].clone())
        self.A_imag = nn.Parameter(A[1].clone())

    def update_state(self, h, x_t):
        # h: [B, 2, H], x_t: [B, H]
        hr, hi = h[:, 0, :], h[:, 1, :]  # [B, H]
        hrn = hr * self.A_real - hi * self.A_imag + x_t  # [B, H]
        hin = hi * self.A_real + hr * self.A_imag        # [B, H]
        hn = torch.stack([hrn, hin], dim=1)              # [B, 2, H]
        return hn, hrn

    def forward(self, u):
        # u: [B, 1, T]
        x = torch.matmul(self.B, u)  # (B, H, T)
        B, H, T = x.shape

        h = torch.zeros(B, 2, H, device=x.device, dtype=x.dtype)  # [B, 2, H]
        h_accum = torch.zeros(B, H, T, device=x.device, dtype=x.dtype)  # [B, H, T]
        i = torch.tensor(0, device=x.device, dtype=torch.int64)
        one = torch.tensor(1, device=x.device, dtype=torch.int64)

        def cond(i, h, h_accum):
            return i < T

        def body(i, h, h_accum):
            x_t = x.index_select(-1, i.unsqueeze(0)).squeeze(-1)  # ✅ safe for export
            h, hr = self.update_state(h, x_t)  # h: [B, 2, H], hr: [B, H]
            h_accum = h_accum.index_copy(-1, i.unsqueeze(0), hr.unsqueeze(-1))  # [B, H, T]
            i_next = i + one
            return i_next, h, h_accum

        _, h, h_accum = torch._higher_order_ops.while_loop(cond, body, (i, h, h_accum))
        y = torch.matmul(self.C, h_accum).transpose(0, 1)  # (B, C_out, T)
        return y

# Instantiate and export
model = ExportableLoop(hidden_size=128, out_channels=10)
inp = torch.randn(1, 1, 32)  # (B, in_channels=1, T=32)
exported = export(model, (inp,))
print("Exporting Done...")
executorch_program = to_edge_transform_and_lower(
    exported
)
print("Edge transform Done...")

executorch_program = executorch_program.to_executorch()  ################Added this line

out_path = "loop_model.pte"
with open(out_path, "wb") as file:
    file.write(executorch_program.buffer)

print(f"Succesfully saved model as {out_path}")

Note:
The support for while_loop has been recently added as part of PR#12062.
For more info see this ticket #8769_comment

Alternatives

No response

Additional context

No response

RFC (Optional)

No response

cc @JacobSzwejbka @angelayi

[vikasbalaga]

In the Core Aten operator set, I can see the aten.squeeze.dim and aten.select are available (not sure if they are compatible to my case or not)

[digantdesai]

We should have both (at least out variants) - https://github.com/pytorch/executorch/tree/main/kernels/portable/cpu.

RuntimeError: Missing out variants: {'aten::squeeze', 'aten::select'}

Can you please share the exported graph after export, to_edge, and to_executorch to help me better understand the issue? Thanks.

[vikasbalaga]

@digantdesai, thanks for your response,
I have shared my script which will perform export(), to_edge() and to_executorch() and the issue is reproducible.

Is that sufficient?

[vikasbalaga]

Hi @digantdesai

Any luck in debugging this issue?

If the script mentioned in the description is not sufficient, I can share the exported graph info that you asked. (Although I am not sure if it means to just print the model object and share the logs after each step?)

[JacobSzwejbka]

Both of those ops are aliasing. ET doesnt support them today. They have to become squeeze_copy and select_copy.

This should be handled by https://github.com/pytorch/executorch/blob/main/exir/passes/replace_broken_ops_with_function_ops_pass.py#L29

But it doesn't seem like select.int is in the list nor does it seem like its working at all with while since unsqueeze is in the list.

[JacobSzwejbka]

Ah ok 1. its squeeze not unsqueeze. and 2 that pass only handles squeeze.default not squeeze.dim.

But tldr its because they havent been converted to copy ops

[vikasbalaga]

@JacobSzwejbka , is there a workaround for this, as I am being blocked by it to convert the model and test inference in ExecuTorch?

Both of those ops are aliasing. ET doesnt support them today. They have to become squeeze_copy and select_copy

Does that mean support needs to be added? If that is the case, I can be helpful in contributing to it (if there are any references, that would be helpful)

I need to figure out a way to convert the model to *.pte and try inference :)

[vikasbalaga]

Hi, any update on this?