nn.RNN — tensor shapes through the model

Background

With batch_first=True, a batch of sequences enters nn.RNN as (B, T, F) and the tensors flow:

$$x:(B, T, F) \;\to\; \text{RNN} \;\to\; \text{out}:(B, T, H) \;\to\; \text{fc(last)} \;\to\; (B, O)$$

and the initial hidden state has shape (num_layers, B, H).

Problem statement

Implement rnn_shapes(B, T, F, H, num_layers, O) returning the key tensor shapes as a dict.

Input

• B, T, F — batch size, sequence length, feature size.
• H — hidden size; num_layers — stacked RNN layers; O — output size.

Output

Returns a dict with keys:

• "input" → (B, T, F)
• "out" → (B, T, H)
• "h0" → (num_layers, B, H)
• "prediction" → (B, O)

Examples

Example 1 — the StockRNN config

Input:  B=32, T=3, F=1, H=64, num_layers=2, O=1
Output: {"input": (32,3,1), "out": (32,3,64), "h0": (2,32,64), "prediction": (32,1)}

Constraints

• Use the shapes above exactly (tuples).
• Return a dict with those four keys.

Notes

• out keeps every timestep's hidden vector; the many-to-one head reads only the last one, collapsing (B, T, H) to (B, H) before the linear layer.