Stop storing None in recurrent dropout state

[rootkiller6788]

Fixes #161.

snt.dynamic_unroll(snt.lstm_with_recurrent_dropout(...), inputs, initial_state) works in eager mode but crashes with ValueError: None values not supported as soon as it's called inside a tf.function. The dropout wrapper's initial_state used None for state elements that aren't dropped (the cell), and dynamic_unroll carries that state through a tf.while_loop, which can't handle None leaves.

I changed maybe_dropout to return an all-ones mask instead of None for elements that shouldn't be dropped. The state stays a plain nest of tensors, and since s * ones == s the dropout behaviour is identical. The existing eager tests still pass, and I added a regression test that runs dynamic_unroll under tf.function (which is what actually exercises the while-loop path) and checks initial_state has no None leaves.

Verified against tensorflow 2.12: the new test fails on the old code and passes with this change; the full recurrent_test.py run is green (94 passed, 15 skipped on CPU-only for GPU/TPU-specific tests).

[sylvesterkaczmarek]

Replacing the undropped cell mask with tf.ones_like(s) fixes the None leaf, but it also adds a full [batch, hidden] tensor to the recurrent state and carries/multiplies it on every tf.while_loop iteration solely to represent identity. For large LSTMs this is avoidable state and memory traffic. A scalar one tensor of the same dtype (broadcast on multiply), or another tensor sentinel, satisfies the loop-state requirement without duplicating the cell-state shape. Please keep the no-None fix without introducing a full-size identity mask.

[rootkiller6788]

Thanks @sylvesterkaczmarek — agreed, and this is addressed in the commits I just pushed. maybe_dropout now returns tf.ones([], dtype=s.dtype) for undropped elements instead of tf.ones_like(s), so the identity mask is a scalar of the same dtype that broadcasts on multiply. It satisfies the loop-state requirement without adding a [batch, hidden] tensor to the recurrent state, and the state still has no None leaves. I also extended the regression test to pin the shape — assertIn((), [tuple(m.shape) for m in tree.flatten(dropout_masks)]) — so a full-size mask can't creep back in. Re-verified on the updated head with TensorFlow 2.16.2 / Python 3.11 (CPU): the issue #161 repro fails on v2 with ValueError: None values not supported and passes with the fix, and the full recurrent_test.py is 94 passed, 15 skipped.

[sylvesterkaczmarek]

Rechecked feb31c0. Undropped recurrent elements now use a scalar one of the cell-state dtype, so the tf.function loop state stays free of None without carrying a full [batch, hidden] identity tensor. The regression pins the scalar shape. My concern is resolved.