Converting the fold operator resulted in an error: Only the Update, Add, Multiply, Maximum, and Minimum operations are currently supported for stablehlo.scatter. Node number 7 (STABLEHLO_SCATTER) failed to prepare

[fiberflow]

Description of the bug:

import torch
import torch.nn as nn
class Fold(nn.Module):
    def __init__(self):
        super(Fold, self).__init__()
        self.fold = nn.Fold(output_size=(4, 5), kernel_size=(2, 2))
    def forward(self, x):
        return self.fold(x)
fold = Fold()
sample = torch.rand(1, 3 * 2 * 2, 12)
print(fold(sample))
fold_edge = ai_edge_torch.convert(fold.eval(), (sample,))
fold_edge(sample)

Actual vs expected behavior:

No response

Any other information you'd like to share?

No response

[pkgoogle]

I was able to replicate w/ the latest commit: 29aa17f

import torch
import torch.nn as nn
import ai_edge_torch
from ai_edge_torch.debug import find_culprits

        
fold = nn.Fold(output_size=(4, 5), kernel_size=(2, 2))
sample = torch.rand(1, 3 * 2 * 2, 12)

print(fold(sample))

fold_edge = ai_edge_torch.convert(fold.eval(), (sample,))
fold_edge(sample)

produces:

Traceback (most recent call last):
  File "xxxxxx/issues/ai-edge-torch/337/test.py", line 13, in <module>
    fold_edge(sample)
  File "xxxxxx/git/ai-edge-torch/ai_edge_torch/model.py", line 106, in __call__
    interpreter.allocate_tensors()
  File "xxxxxx/envs/aet_head/lib/python3.11/site-packages/ai_edge_litert/interpreter.py", line 534, in allocate_tensors
    return self._interpreter.AllocateTensors()
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
RuntimeError: Only update, Add, Multiply, Maximum and Minimum operations are currently supported for stablehlo.scatter.Node number 7 (STABLEHLO_SCATTER) failed to prepare.

find_culprits produces a different error:

culprits = find_culprits(fold.eval(), (sample,))
culprit = next(culprits)
culprit.print_code()

Traceback (most recent call last):
  File "xxxxxx/issues/ai-edge-torch/337/test.py", line 15, in <module>
    culprit = next(culprits)
              ^^^^^^^^^^^^^^
  File "xxxxxx/git/ai-edge-torch/ai_edge_torch/debug/culprit.py", line 485, in find_culprits
    for search_result in _search_model(
  File "xxxxxx/git/ai-edge-torch/ai_edge_torch/debug/culprit.py", line 457, in _search_model
    raise e
  File "xxxxxx/git/ai-edge-torch/ai_edge_torch/debug/culprit.py", line 434, in _search_model
    raw_min_fx_gm, raw_min_inputs = fx_minifier(
                                    ^^^^^^^^^^^^
  File "xxxxxx/envs/aet_head/lib/python3.11/site-packages/torch/_functorch/fx_minifier.py", line 227, in minifier
    raise RuntimeError("Input graph did not fail the tester")
RuntimeError: Input graph did not fail the tester

I should note the conversion "works" but it fails when attempting to run/be called.