Add Bidirectional #2770
Add Bidirectional #2770
Conversation
cgarciae
force-pushed
the
add-bidirectional
branch
from
dd4fdf5 to
b5d50c5
Compare
cgarciae
force-pushed
the
add-bidirectional
branch
from
c51c9f6 to
1834bdb
Compare
cgarciae
force-pushed
the
add-bidirectional
branch
from
c0314ad to
da12439
Compare
Codecov Report
@@ Coverage Diff @@
## main #2770 +/- ##
==========================================
+ Coverage 81.91% 82.01% +0.09%
==========================================
Files 55 55
Lines 5918 5977 +59
==========================================
+ Hits 4848 4902 +54
- Misses 1070 1075 +5
Help us with your feedback. Take ten seconds to tell us how you rate us. Have a feature suggestion? Share it here. |
| kernel = jnp.concatenate(kernels, axis=-1) | ||
| if use_bias: | ||
| biases = [] |
There was a problem hiding this comment.
This code is created to make mypy happy but also improves the correctness of the program.
| @@ -717,6 +744,12 @@ def scan_fn( | |||
| else: | |||
| carry, outputs = scan_output | |||
|
|
|||
| if reverse and keep_order: | |||
| outputs = jax.tree_map( | |||
There was a problem hiding this comment.
maybe factor this out into a separate function since you use it twice?
There was a problem hiding this comment.
Are we DRY or WET? For this use case I think abstracting away a tree_map leads to unnecessary indirection.
cgarciae
force-pushed
the
add-bidirectional
branch
2 times, most recently
from
b7c0a91 to
984ec14
Compare
flax/linen/recurrent.py
Outdated
| @@ -651,16 +664,27 @@ def __call__( | |||
| else it will return a tuple of the final carry and the output sequence. | |||
| time_major: if ``time_major=False`` (default) it will expect inputs with shape | |||
| ``(*batch, time, *features)``, else it will expect inputs with shape ``(time, *batch, *features)``. | |||
| reverse: overrides the ``reverse`` attribute, if ``reverse=False`` (default) the sequence | |||
| is processed from left to right, else it will be processed from right to left, output order | |||
| will be the same as input order. If ``segmentation_mask`` is passed, padding elements will | |||
There was a problem hiding this comment.
I think keep_order is still a bit difficult to understand, so good documentation is important. This might be too unclear:
reverse: overrides the ``reverse`` attribute, if ``reverse=False`` (default) the sequence
is processed from left to right, else it will be processed from right to left, output order
will be the same as input order.
what about:
reverse: overrides the reverse attribute, if reverse=False (default) the sequence
is processed from left to right and returned in the original order, else it will be processed from right to left, and returned in reverse order (but with the padding values at the end).
maybe even a simple example could be good. e.g.
`[1, 2, 3, 0, 0]` vs `[3, 2, 1, 0, 0]`?
984ec14 to
74e00f2
Compare
cgarciae
force-pushed
the
add-bidirectional
branch
from
74e00f2 to
6412c6c
Compare
What does this PR do?
BidirectionalModule that takes twoRNNinstances and applies the bidirectional logic.reverseandkeep_orderarguments toRNNto allow the implementation of the bidirectional logic.ArrayfromAnytojax.Arrayand fixes some type errors.