[Pallas][MPMD] Preserve out_avals memory spaces in _mpmd_map_discharge_rule #40777

[Pallas][MPMD] Preserve out_avals memory spaces in _mpmd_map_discharge_rule

_mpmd_map_abstract_eval strips Pallas memory spaces from eqn.outvars to
jax_core.MemorySpace.Device while eqn.params["out_avals"] retains the
original Pallas memory spaces. Previously, _mpmd_map_discharge_rule discarded
eqn.params["out_avals"] and constructed updated_out_avals from
ctx.out_avals, erasing out_type memory spaces on discharged mpmd_map
kernels and omitting output_memory_space_colors during TPU lowering.

Preserve out_avals in _mpmd_map_discharge_rule and add a unit test in
tpu_pallas_mpmd_test.py.