TorchScript: Unsupported 'torch.polar' expression type

[hzhangxyz]

TorchScript: Unsupported 'torch.polar' expression type

[lutzroeder]

@hzhangxyz can you share a sample file to reproduce the issue.

[hzhangxyz]

@lutzroeder

Here is an example to reproduce it:

import torch


class Example(torch.nn.Module):

    def forward(self, x):
        return torch.polar(torch.ones_like(x), x)


example = torch.jit.script(Example())

print(example(torch.arange(10, dtype=torch.float32)))

torch.jit.save(example, "example.pt")

[hzhangxyz]

torch.polar is used in llama to implement rotary positional embedding, see https://github.com/meta-llama/llama/blob/54c22c0d63a3f3c9e77f43a6a3041c00018f4964/llama/model.py#L103