Open domschl opened 2 weeks ago
Hi @domschl, thanks for the bug report. The bug is from the compilation. There is a subgraph that is too big (actually just has too many inputs) to fuse into a single kernel but compile still tries and fails. This big cryptic string is actually a representation of the graph to be fused.
We 'll look into fixing it (ie compile should break the subgraph in two smaller ones). In the meantime, you could disable compile and the code should run fine.
Tx! Confirmed: without compilation it works fine.
Describe the bug
In:
If self.context_recurrent is any of nn.LSTM, RNN, or GRU, the above crash happens.
To Reproduce
Insert
x = nn.RNN(x)
between two transformer layers causes the crash.Complete code https://github.com/domschl/mlx-poet/blob/cedac548256a1bd2a1bb33362cf9d99f22a360c7/mlx_poet_bug.py (requires
pip install ml-indie-tools
)Expected behavior
No crash, and if necessary clear error message. I've checked that there is no tensor-shape problem.
Desktop (please complete the following information):
Additional context Add any other context about the problem here.