LiSsHhUuAaIi opened a new issue, #18477:
URL: https://github.com/apache/tvm/issues/18477
### Description
When converting a PyTorch model containing Gumbel Softmax operations and
related stochastic operations to TVM Relax module via `torch.export`, an
AssertionError occurs. TVM currently does not support the
`exponential_.default`, `scatter_.value`, and `max.dim` operations that are
essential for Gumbel Softmax and stochastic masking.
### Expected behavior
The PyTorch model with Gumbel Softmax and stochastic operations should be
successfully converted to TVM Relax module, enabling deployment of models that
use categorical reparameterization and stochastic components.
### Actual behavior
An AssertionError occurs during `from_exported_program` conversion with the
message `Unsupported function types ['exponential_.default', 'max.dim',
'scatter_.value']`, indicating that TVM's PyTorch frontend lacks support for
these operations.
### Environment
- OS: Ubuntu 20.04.6 LTS
- TVM version: 0.23.dev0
- Python version: 3.11.14
### Steps to reproduce
```python
import torch
import torch.nn as nn
import torch.nn.functional as F
import tvm
from tvm import relax
class MinimalGumbelModel(nn.Module):
def __init__(self):
super(MinimalGumbelModel, self).__init__()
def forward(self, x):
# Gumbel Softmax operation (contains multiple unsupported ops)
x = F.gumbel_softmax(x, tau=1.0, hard=True)
return x
model = MinimalGumbelModel()
model.eval()
x = torch.randn(32, 10)
# PyTorch execution works
with torch.no_grad():
output = model(x)
# PyTorch export works
exported_program = torch.export.export(model, (x,))
# TVM conversion fails
from tvm.relax.frontend.torch import from_exported_program
mod = from_exported_program(exported_program) # AssertionError here
```
### Triage
* needs-triage
* frontend: pytorch
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
