Combine into a TabPFN-style two-way attention block

The problem, from first principles

01-row-wise-attention and 02-column-wise-attention are two separate halves of the same idea — one axis alone can't capture everything a table has to say about itself. TabPFN's actual architecture alternates these two directions, letting information mix across an entire row, then across an entire column, then row again, and so on, building up genuinely table-wide context: after enough layers, information from any cell can eventually influence any other cell.

From theory to code

Implement two_way_attention_block(table, row_weights, col_weights): row_wise_attention, then column_wise_attention on its output, each with a residual connection. Reuse 01-row-wise-attention's row_wise_attention and 02-column-wise-attention's column_wise_attention. Each attention step adds its own input back to its output (a residual connection), not just the raw attention result.

Constraints

• table: shape (n_rows, n_cols, d_model). Returns the same shape.
• row_weights/col_weights: each a (w_query, w_key, w_value) tuple.
• Column-wise attention operates on the row-wise step's output, not the original table.
• With all-zero weight matrices, the output must equal the input exactly (both attention steps contribute nothing, leaving only the residual passthrough).

Hints

Open one at a time. Each gives away a little more than the last.