Optimize RoPE arithmetic and simplify tensor reshaping

[Perseus14]

This PR introduces two focused optimizations to improve memory bandwidth utilization and overall code maintainability. It optimizes the Rotary Position Embedding (RoPE) application by utilizing native real arithmetic for better hardware fusion, and refactors a sequence of tensor collapses into a single, robust reshape operation.

Changes:

• Replaced the float32 upcasting and jax.lax.complex multiplication with explicit real-number arithmetic (computing the 2D rotation directly), keeping the tensors in their native dtype (e.g., bfloat16).

• Replaced three sequential jax.lax.collapse operations with a single, explicit hidden_states.reshape(batch_size, -1, num_frames, height, width).

[github-actions]

e2e testgrid: https://8bcf50593faf4ea38060e236169827e5-dot-us-central1.composer.googleusercontent.com/dags/maxdiffusion_tpu_e2e/grid

[mbohlool]

LGTM