Incompatible with PyTorch 2.0's torch.compile()

[zaptrem]

Using rearrange_many in a model results in a very nondescript error when running torch.compile(model)

However, replacing rearrange_many with rearrange fixes it. Could it be related to your use of wrapper/higher-order functions?

Said non-descript error:

Traceback (most recent call last):
  File "/home/zaptrem/dance3/sonauto-models/env/lib/python3.9/site-packages/torch/_dynamo/convert_frame.py", line 398, in _compile
    out_code = transform_code_object(code, transform)
  File "/home/zaptrem/dance3/sonauto-models/env/lib/python3.9/site-packages/torch/_dynamo/bytecode_transformation.py", line 341, in transform_code_object
    transformations(instructions, code_options)
  File "/home/zaptrem/dance3/sonauto-models/env/lib/python3.9/site-packages/torch/_dynamo/convert_frame.py", line 385, in transform
    tracer.run()
  File "/home/zaptrem/dance3/sonauto-models/env/lib/python3.9/site-packages/torch/_dynamo/symbolic_convert.py", line 1676, in run
    super().run()
  File "/home/zaptrem/dance3/sonauto-models/env/lib/python3.9/site-packages/torch/_dynamo/symbolic_convert.py", line 528, in run
    and self.step()
  File "/home/zaptrem/dance3/sonauto-models/env/lib/python3.9/site-packages/torch/_dynamo/symbolic_convert.py", line 496, in step
    getattr(self, inst.opname)(inst)
  File "/home/zaptrem/dance3/sonauto-models/env/lib/python3.9/site-packages/torch/_dynamo/symbolic_convert.py", line 305, in wrapper
    return inner_fn(self, inst)
  File "/home/zaptrem/dance3/sonauto-models/env/lib/python3.9/site-packages/torch/_dynamo/symbolic_convert.py", line 956, in CALL_FUNCTION
    self.call_function(fn, args, {})
  File "/home/zaptrem/dance3/sonauto-models/env/lib/python3.9/site-packages/torch/_dynamo/symbolic_convert.py", line 430, in call_function
    self.push(fn.call_function(self, args, kwargs))
  File "/home/zaptrem/dance3/sonauto-models/env/lib/python3.9/site-packages/torch/_dynamo/variables/nn_module.py", line 220, in call_function
    return tx.inline_user_function_return(
  File "/home/zaptrem/dance3/sonauto-models/env/lib/python3.9/site-packages/torch/_dynamo/symbolic_convert.py", line 466, in inline_user_function_return
    result = InliningInstructionTranslator.inline_call(self, fn, args, kwargs)
  File "/home/zaptrem/dance3/sonauto-models/env/lib/python3.9/site-packages/torch/_dynamo/symbolic_convert.py", line 1750, in inline_call
    return cls.inline_call_(parent, func, args, kwargs)
  File "/home/zaptrem/dance3/sonauto-models/env/lib/python3.9/site-packages/torch/_dynamo/symbolic_convert.py", line 1805, in inline_call_
    tracer.run()
  File "/home/zaptrem/dance3/sonauto-models/env/lib/python3.9/site-packages/torch/_dynamo/symbolic_convert.py", line 528, in run
    and self.step()
  File "/home/zaptrem/dance3/sonauto-models/env/lib/python3.9/site-packages/torch/_dynamo/symbolic_convert.py", line 496, in step
    getattr(self, inst.opname)(inst)
  File "/home/zaptrem/dance3/sonauto-models/env/lib/python3.9/site-packages/torch/_dynamo/symbolic_convert.py", line 305, in wrapper
    return inner_fn(self, inst)
  File "/home/zaptrem/dance3/sonauto-models/env/lib/python3.9/site-packages/torch/_dynamo/symbolic_convert.py", line 1005, in CALL_FUNCTION_KW
    self.call_function(fn, args, kwargs)
  File "/home/zaptrem/dance3/sonauto-models/env/lib/python3.9/site-packages/torch/_dynamo/symbolic_convert.py", line 430, in call_function
    self.push(fn.call_function(self, args, kwargs))
  File "/home/zaptrem/dance3/sonauto-models/env/lib/python3.9/site-packages/torch/_dynamo/variables/functions.py", line 214, in call_function
    return super(UserFunctionVariable, self).call_function(tx, args, kwargs)
  File "/home/zaptrem/dance3/sonauto-models/env/lib/python3.9/site-packages/torch/_dynamo/variables/functions.py", line 67, in call_function
    return tx.inline_user_function_return(
  File "/home/zaptrem/dance3/sonauto-models/env/lib/python3.9/site-packages/torch/_dynamo/symbolic_convert.py", line 466, in inline_user_function_return
    result = InliningInstructionTranslator.inline_call(self, fn, args, kwargs)
  File "/home/zaptrem/dance3/sonauto-models/env/lib/python3.9/site-packages/torch/_dynamo/symbolic_convert.py", line 1750, in inline_call
    return cls.inline_call_(parent, func, args, kwargs)
  File "/home/zaptrem/dance3/sonauto-models/env/lib/python3.9/site-packages/torch/_dynamo/symbolic_convert.py", line 1805, in inline_call_
    tracer.run()
  File "/home/zaptrem/dance3/sonauto-models/env/lib/python3.9/site-packages/torch/_dynamo/symbolic_convert.py", line 528, in run
    and self.step()
  File "/home/zaptrem/dance3/sonauto-models/env/lib/python3.9/site-packages/torch/_dynamo/symbolic_convert.py", line 496, in step
    getattr(self, inst.opname)(inst)
  File "/home/zaptrem/dance3/sonauto-models/env/lib/python3.9/site-packages/torch/_dynamo/symbolic_convert.py", line 1911, in LOAD_CLOSURE
    self.push(self.closure_cells[inst.argval])
KeyError: 'fn\n\nfrom user code:\n   File "/home/zaptrem/dance3/audio-diffusion-pytorch/audio_diffusion_pytorch/modules.py", line 354, in forward\n    q, k, v = rearrange_many((q, k, v), "b n (h d) -> b h n d", h=self.num_heads)\n  File "/home/zaptrem/dance3/sonauto-models/env/lib/python3.9/site-packages/einops_exts/einops_exts.py", line 19, in inner\n    return (fn(tensor, pattern, **kwargs) for tensor in tensors)\n\n\nYou can suppress this exception and fall back to eager by setting:\n    torch._dynamo.config.suppress_errors = True\n'