How do you set up a stateful onnx model.

[JamesBowerXanda]

Hi,

I am trying to set up a simple stateful model in triton inference server to get used to the syntax. I would like to use a staeful onnx model with an RNN like structure. Here is some simple code to create the onnx file:

from torch import nn
import torch

class StateFullModule(nn.Module):
    def __init__(self):
        super(StateFullModule, self).__init__()
        self.linear_hidden = nn.Linear(5, 10)
        self.linear_out = nn.Linear(10, 5)

    def forward(self, input_token, state):
        hidden = self.linear_hidden(input_token)
        new_state = hidden + state
        new_state = nn.functional.normalize(new_state, dim=1)
        
        # Explicitly handle the shape with view
        batch_size = new_state.size(0)
        new_state = new_state.view(batch_size, 10)
        
        out = self.linear_out(new_state)
        out = nn.functional.softmax(out, dim=1)
        
        return out, new_state
    
model = StateFullModule()

# Creating dummy input tensors with correct shapes
input_token = torch.randn(1, 5)  # Batch size 1, input dimension 5
past_state = torch.randn(1, 10)  # Batch size 1, hidden state dimension 10

# Export the model to ONNX format
torch.onnx.export(
    model,
    (input_token, past_state),
    "model.onnx",
    input_names=["input_token", "past_state"],
    output_names=["output", "present_state"],
    dynamic_axes={
        "input_token": {0: "batch_size"},
        "past_state": {0: "batch_size"},
        "output": {0: "batch_size"},
        "present_state": {0: "batch_size"}
    },
    opset_version=13
)

I would like to call the model from a python BLS model. For now lets say we will limit the batch_size to 1 and the input to the python model will be an integer called "iterations" which just says how many times to run the rnn. What do the configs and model.py file look like? I cannot find any examples online.