NsNet2 sample fails with unexpected 'constraints' argument error

[mig9mili]

## Description
When building and testing Quidditch, the NsNet2 sample fails with an error indicating an unexpected `constraints` argument being passed to `NsNet2.forward()`.

## Error Message
```python
TypeError: NsNet2.forward() got an unexpected keyword argument 'constraints'

Steps to Reproduce

1. Clone the repository
2. Set up environment with required dependencies
3. Run either:

   cmake --build . --target test
   # or
   cmake --build . --target codegen runtime

Full Error Stack Trace

Traceback (most recent call last):
  File ".../NsNet2.py", line 102, in <module>
    exported = aot.export(with_frames(n_frames=args.frames))
  File ".../shark_turbine/aot/exporter.py", line 321, in export
    cm = TransformedModule(context=context, import_to="import")
  [...]
  File ".../shark_turbine/aot/builtins/jittable.py", line 197, in flat_wrapped_f
    return self.wrapped_f(*pytorch_args, **pytorch_kwargs)
TypeError: NsNet2.forward() got an unexpected keyword argument 'constraints'

Environment

• OS: Ubuntu
• Python Version: 3.11.11
• CMake Version: 3.22.1
• Using virtual environment with required dependencies installed

Additional Context

The error seems to occur during the IREE compilation pipeline when trying to export the NsNet2 model. The error suggests there's a mismatch between the model's forward method signature and how it's being called during the compilation process.

Would you like me to explain any part of this issue template or help modify it further?

[superlopuh]

CC @zero9178 @compor

[superlopuh]

At first glance this seems to be a shark_turbine error, and not quidditch-specific, but that probably means that Quidditch should pin dependencies more precisely if this worked previously.

[Groverkss]

Shark-Turbine was moved to https://github.com/iree-org/iree-turbine some time ago. So i'd say quidditch is running behind by a lot. You would need to explicitly pin a older version of pytorch + shark_turbine. Newer versions of pytorch use a different API.

[mig9mili]

if you can let me know which version worked previously ? it would be great

[compor]

Hi there @mig9mili it was iree-turbine 2.3.0, also noted under runtime/samples/nsnet2/requirements.txt

This is what I'm using ATM, part of my pip list

Package                   Version             Editable project location
------------------------- ------------------- ------------------------------
attrs                     24.2.0
filelock                  3.16.1
fsspec                    2024.9.0
hjson                     3.1.0
immutabledict             4.2.0
iree-compiler             20240828.999
iree-runtime              20240828.999
iree-turbine              2.3.0
...

Can you confirm this on your side too?

[compor]

Shark-Turbine was moved to https://github.com/iree-org/iree-turbine some time ago. So i'd say quidditch is running behind by a lot. You would need to explicitly pin a older version of pytorch + shark_turbine. Newer versions of pytorch use a different API.

@Groverkss how fussy and desirable would it be to update this? any pointers appreciated.

[compor]

I'll try to pin this dependency more precisely, as I've bumped into this myself on a new environment