flax

Add test case for degenerate and deterministic attention layer dropout.

#3083

Merged