[Feature Request] Support for Gumbel Softmax and related stochastic operations in PyTorch frontend

[LiSsHhUuAaIi]

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

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