Fix the GPUA version of MRG to use the same rstate shape as the CPU one.