Custom TypeVar doesn't work if it's inside a function

[titaiwangms]

Found this is testing,

Works:

TCustomFloat = TypeVar("TCustomFloat", bound=Union[FLOAT16, FLOAT, DOUBLE, BFLOAT16])

@onnxscript.script(custom_domain)
def test_custom_op(x: TCustomFloat, y: TCustomFloat) -> TCustomFloat:
    return op.Add(x, y)

No:

def test_something(self):
  TCustomFloat = TypeVar("TCustomFloat", bound=Union[FLOAT16, FLOAT, DOUBLE, BFLOAT16])
  
  @onnxscript.script(custom_domain)
  def test_custom_op(x: TCustomFloat, y: TCustomFloat) -> TCustomFloat:
      return op.Add(x, y)

[justinchuby]

There may be scoping issues in the translation process we need to fix

[gramalingam]

Correct, this is a known issue. This code needs to be updated. But it is a bit non-trivial. We need to be able to access the non-local environment at the definition-site at run-time. Suggestions on how to do this welcome.