[Refactor] Update type references from tir to tirx in PyTorch ExportedProgram frontend

[mshr-h]

Follow up for #18913 and #18917

[gemini-code-assist]

Summary of Changes

Hello, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request refactors the exported_program_translator.py file by updating its dependencies from the tir module to the tirx module. This change is part of a broader effort to standardize module naming and usage, ensuring consistency across the codebase and aligning with recent architectural updates.

Highlights

• Refactor tir to tirx imports: The pull request updates all references from the tir module to tirx within the exported_program_translator.py file. This includes both the top-level import statement and subsequent usage of tir.PrimExpr and tir.analysis.expr_deep_equal.

🧠 New Feature in Public Preview: You can now enable Memory to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console.

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point by creating a comment using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

Feature	Command	Description
Code Review	/gemini review	Performs a code review for the current pull request in its current state.
Pull Request Summary	/gemini summary	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in pull request comments and review comments.
Help	/gemini help	Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for GitHub and other Google products, sign up here.

Footnotes

1. Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

[gemini-code-assist]

Code Review

This pull request refactors python/tvm/relax/frontend/torch/exported_program_translator.py to use the tvm.tirx namespace instead of tvm.tir. The changes correctly update the import statement and all usages of tir objects and functions to their tirx counterparts. The refactoring is consistent with the ongoing migration to the new tirx namespace and the changes appear correct.

[mshr-h]

Now we can import and execute Whisper from OpenAI.
Only support static shape input for now.

whisper.py

import numpy as np
import torch
from torch.export import export
from transformers import AutoProcessor, WhisperForConditionalGeneration
from tvm.relax.frontend.torch import from_exported_program

import tvm
from tvm import relax

MODEL_ID = "openai/whisper-tiny"

processor = AutoProcessor.from_pretrained(MODEL_ID)
hf_model = WhisperForConditionalGeneration.from_pretrained(MODEL_ID).eval()

audio = np.load("audio_16khz_mono.npy").astype(np.float32)

inputs = processor(audio, sampling_rate=16000, return_tensors="pt")
input_features = inputs.input_features  # [B, 80, T]

with torch.no_grad():
    hf_generated = hf_model.generate(input_features, max_new_tokens=128)
hf_text = processor.batch_decode(hf_generated, skip_special_tokens=True)[0]

# Fixed max decoder length for static-shape compilation
MAX_DEC_LEN = 128
PAD_TOKEN_ID = (
    hf_model.config.pad_token_id if hf_model.config.pad_token_id is not None else 0
)


class WhisperNoCache(torch.nn.Module):
    def __init__(self, model):
        super().__init__()
        self.model = model

    def forward(self, input_features, decoder_input_ids):
        out = self.model(
            input_features=input_features,
            decoder_input_ids=decoder_input_ids,
            use_cache=False,
            return_dict=False,
        )
        return out[0]  # logits


wrapped = WhisperNoCache(hf_model)

# Trace with fixed-size decoder_input_ids (padded to MAX_DEC_LEN)
decoder_ids_trace = torch.full((1, MAX_DEC_LEN), PAD_TOKEN_ID, dtype=torch.long)
decoder_ids_trace[0, 0] = hf_model.config.decoder_start_token_id

with torch.no_grad():
    exported_program = export(wrapped, (input_features, decoder_ids_trace))
    mod = from_exported_program(exported_program, keep_params_as_input=True)
    mod, params = relax.frontend.detach_params(mod)

dev = tvm.cuda(0) if tvm.cuda(0).exist else tvm.cpu(0)
target = tvm.target.Target.from_device(dev)

s_tir_pipeline = tvm.transform.Sequential(
    [
        tvm.s_tir.transform.DefaultGPUSchedule(),
        tvm.s_tir.pipeline.default_s_tir_pipeline(),
    ]
)

ex = tvm.compile(mod, target=target, tir_pipeline=s_tir_pipeline)
vm = relax.VirtualMachine(ex, dev)

params_tvm = [tvm.runtime.tensor(p, dev) for p in params["main"]]
features_tvm = tvm.runtime.tensor(input_features.numpy(), dev)

# Autoregressive decoding with fixed-length padded input
decoder_ids = torch.full((1, MAX_DEC_LEN), PAD_TOKEN_ID, dtype=torch.long)
decoder_ids[0, 0] = hf_model.config.decoder_start_token_id
cur_len = 1

for _ in range(MAX_DEC_LEN - 1):
    dec_tvm = tvm.runtime.tensor(decoder_ids.numpy(), dev)
    logits = vm["main"](features_tvm, dec_tvm, *params_tvm)

    if not hasattr(logits, "numpy"):
        logits = logits[0]

    # Take logits at current position (last non-pad token)
    next_id = int(logits.numpy()[0, cur_len - 1].argmax())
    decoder_ids[0, cur_len] = next_id
    cur_len += 1

    if next_id == hf_model.config.eos_token_id:
        break

generated = decoder_ids[:, :cur_len]
tvm_text = processor.batch_decode(generated, skip_special_tokens=True)[0]

print(f"[HuggingFace] {hf_text}")
print(f"[TVM]         {tvm_text}")

[mshr-h]

cc @tlopex @guan404ming

[guan404ming]

Cool! Thanks.

[guan404ming]

Now we can import and execute Whisper from OpenAI.

Wow! I didn't think about that before but sounds really interesting. Do you have any demo?

[mshr-h]

Thanks! I don’t have a demo yet, but I’ll share updates once we have a prototype.

[guan404ming]

Cool! Looking forward to that.