[ROCm] Implement RNN support#25755

copybara-service[bot] merged 1 commit intojax-ml:mainjax-ml/jax:mainfrom

ROCm:ci_rnn_final-upstreamROCm/jax:ci_rnn_final-upstreamCopy head branch name to clipboard