Read the packed-sequence schema of gru / rnn_tanh / rnn_relu correctly

[LeSingh1]

Problem

Feeding a PackedSequence to nn.RNN or nn.GRU aborts conversion:

import torch, torch.nn as nn, coremltools as ct
from torch.nn.utils.rnn import pack_padded_sequence, pad_packed_sequence

class Encoder(nn.Module):
    def __init__(self):
        super().__init__()
        self.rnn = nn.RNN(4, 6, batch_first=True)

    def forward(self, x, lens):
        packed, _ = self.rnn(pack_padded_sequence(x, lens, batch_first=True))
        out, _ = pad_packed_sequence(packed, batch_first=True)
        return out

x, lens = torch.randn(3, 10, 4), torch.tensor([10, 5, 1], dtype=torch.int32)
ct.convert(torch.jit.trace(Encoder().eval(), (x, lens)), convert_to="mlprogram", inputs=[...])

ERROR - converting 'rnn_tanh' op
AttributeError: 'list' object has no attribute 'val'

torch has two schemas for each of these ops, and the packed one inserts batch_sizes ahead of hx:

aten::rnn_tanh.input(Tensor input, Tensor hx, Tensor[] params, bool has_biases,
                     int num_layers, float dropout, bool train, bool bidirectional,
                     bool batch_first)
aten::rnn_tanh.data(Tensor data, Tensor batch_sizes, Tensor hx, Tensor[] params,
                    bool has_biases, int num_layers, float dropout, bool train,
                    bool bidirectional)

Both bind exactly 9 inputs, so the expected=9 check passes and nothing catches the shift — every argument from index 1 on is read one slot early. has_bias = inputs[3].val lands on the Tensor[] params list, which is where the AttributeError comes from. gru and rnn_relu have the identical pair of schemas and the identical bug.

lstm already handles this (has_batch_sizes = not isinstance(inputs[1], Iterable)); the single-hidden-state ops just never got the same treatment.

Fix

A shared _parse_rnn_args helper used by gru, rnn_tanh and rnn_relu.

Distinguishing the schemas needs care, because lstm's test does not transfer: lstm's hx is a Tensor[], so a list at index 1 means the unpacked schema — but for these ops hx is a single Tensor, so index 1 is a Var either way. What is unambiguous is the Tensor[] params weight list, the only list-valued argument: it sits at index 2 for .input and index 3 for .data. So isinstance(inputs[2], Iterable) identifies the schema, and everything after is read at the matching offset.

.data carries no batch_first argument, because the output of _pack_padded_sequence is always laid out batch first — same as lstm already assumes.

Test

TestRNNWithPackedSequence::test_rnn runs a packed nn.RNN through pack_padded_sequence / pad_packed_sequence for both nonlinearities (tanh -> aten::rnn_tanh.data, relu -> aten::rnn_relu.data) and both batch_first settings, comparing against torch. It mirrors the existing TestLSTMWithPackedSequence, which is what made the lstm handling visible in the first place.

All 8 parametrizations fail with the AttributeError without the fix and pass with it (max abs difference vs torch ~1.2e-07).

A note on GRU

gru gets the same corrected parsing, but a packed GRU still does not convert end to end — it now fails further downstream with

NotImplementedError: Only static shape of PackedSequence object is supported.

That is a separate, pre-existing limitation: _add_gru_layer builds its hidden-state list with mb.fill over a shape computed at run time, so the GRU output carries a symbolic batch dimension that _pad_packed_sequence rejects. lstm avoids it because _add_mil_lstm keeps a static shape. I left that alone rather than widen this change — the argument-offset fix stands on its own, and it turns a misleading AttributeError deep inside argument parsing into an accurate message about what is actually unsupported. Happy to look at the GRU lowering separately if it is wanted.

Verification

pytest coremltools/converters/mil/frontend/torch/test/test_torch_ops.py -k "TestRNN or TestGRU or TestLSTM"

on macOS / Apple silicon, torch 2.12, mlprogram + neuralnetwork.