Carefully move some internal methods of recurrent.py to a dedicated utils file

[SergiiVolodko]

Before trying to tackle #161 I thought it can be helpful to a bit simplify navigation in recurrent.py by reducing the number of lines of code there. Carefully extracting few internal methods to another file seemed to be a quick win in that sense. So I decided to come up with this PR. Hopefully, it is useful for the maintainability of recurrent.py.

Changes

• Moved following methods from recurrent.py to a new file recurrent_internals.py :

Method
_check_inputs_dtype
_safe_where
_ unstack_input_sequence
_specialize_per_device
_rnn_step
_lstm_fn

• Added recurrent_internals.py as a dependency to the Bazel src/BUILD
• Aligned the usage of recurrent_internals with the usage of utils

The tests are passing locally

Questions

It was quite a useful journey for me and it would be great if you find the outcome useful as well:) I've also got a few questions while implementing the changes.

• [ ] Is recurrent_internals.py a good enough name? Please feel free to suggest a nicer or more consistent one.
• [ ] I’ve noticed quite a few TODOs in recurrent.py. Maybe you have some refactoring plans or ideas for this file which should be considered in the current or next PRs?
• [ ] Are there more places in code that should know about the newly added file?

[googlebot]

We found a Contributor License Agreement for you (the sender of this pull request), but were unable to find agreements for all the commit author(s) or Co-authors. If you authored these, maybe you used a different email address in the git commits than was used to sign the CLA (login here to double check)? If these were authored by someone else, then they will need to sign a CLA as well, and confirm that they're okay with these being contributed to Google. In order to pass this check, please resolve this problem and then comment @googlebot I fixed it.. If the bot doesn't comment, it means it doesn't think anything has changed.

ℹ️ Googlers: Go here for more info.

[SergiiVolodko]

@googlebot I fixed it.

[googlebot]

All (the pull request submitter and all commit authors) CLAs are signed, but one or more commits were authored or co-authored by someone other than the pull request submitter.

We need to confirm that all authors are ok with their commits being contributed to this project. Please have them confirm that by leaving a comment that contains only @googlebot I consent. in this pull request.

Note to project maintainer: There may be cases where the author cannot leave a comment, or the comment is not properly detected as consent. In those cases, you can manually confirm consent of the commit author(s), and set the cla label to yes (if enabled on your project).

ℹ️ Googlers: Go here for more info.

[SergiiVolodko]

@googlebot I consent.

[SergiiVolodkoWorking]

@googlebot I consent.

[googlebot]

CLAs look good, thanks!

ℹ️ Googlers: Go here for more info.

[SergiiVolodko]

Sorry for the mess with the commit authors, a bit an unexpected issue but seems to be solved :) Will squash the commits if needed

[superbobry]

Thanks for the PR, Sergey!

I understand your concerns regarding the size of recurrent.py. Big modules could be tricky to work with, but I'm not sure that splitting recurrent.py into multiple smaller modules would significantly improve readability/maintainability. Having internal functions defined closer to the usage sites makes refactoring easier; _ names clearly indicate that these are not designed for use outside of Sonnet. So, all in all, I'd rather not merge this. I hope it's okay. That said, please do not hesitate to tackle #161, we'd welcome a fix for that!

I’ve noticed quite a few TODOs in recurrent.py. Maybe you have some refactoring plans or ideas for this file which should be considered in the current or next PRs?

I don't think we have any refactoring planned atm. Most TODOs are there to provide context or serve as a reminder, but none are critical to implement/fix.

[SergiiVolodko]

Thank you @superbobry ! Sure, there is no problem to skip the changes. Or maybe to restructure them… I think I was a bit unclear about the reasoning for this PR, even to myself. So let me try to rewrite it :)

The selection of recurrent.py as a refactoring target is not random. When thinking where my contribution could be the most useful I’ve run a great code BI tool against the repository and it revealed that recurrent.py is actually a significant outlier across the entire project. (Legend: Circles are files in folders, their size - number of lines of code, color insensitivity - commit frequency)

It has the highest change frequency and the biggest size.

So the major reason to think of refactoring this file not just the size, but the frequency of changes/fixes. (Because we commit when we need to fix or modify something )

Zooming a bit deeper the analysis identified that there is quite a high chances that static_unroll and dynamic_unroll methods are problematic: ( __init__ and __call__ seems to be aggregated through all the classes so we can skip their results)

And when I was trying to reproduce the example from #161 something was failing in dynamic_unroll for me.

Intuitively, these two methods are actual targets of my refactoring. My intuitive plan was:

1. (current PR) Clean up a bit before the extraction and get familiar with the code and not break anything.
2. Extract static_unroll and dynamic_unroll to a dedicated file like: recurrent_internals.py or unroll.py.

But seems like I completely formulated this plan only now :)

The idea of this refactoring is to extract static_unroll and dynamic_unroll with their internals to a separate file like unroll.py. And I can see 3 major benefits from that:

• Better isolate most frequently changed code, which can help to focus maintenance effort in the future.
• recurrent.py will contain only classes of networks that with time seems to be only getting added, not modified.
• Better align the structure with usage. From the code usages the methods seem to be called like: snt.dynamic_unroll() with no recurrent in the path. While structure-wise, currently, the methods belong to recurrent.py.

Of course, this refactoring idea is based only on numbers without knowing the details of your challenges :)

So if you think the extraction of unroll methods can make sense I would be happy to shift the focus of this PR and implement the change. Otherwise, no problem to close this branch.

How do you think?