At the moment, pytorch version of most attention had been refactored to use attention processor (except a few being deprecated). Is there plans to do the same for flax version?

Since AttentionBlock is deprecated, should the same be done to its flax counterpart?