Add missing perf/loss baselines for paxml

[terrykong]

There are currently still two missing tests for the paxml MGMN test suite:

Tasks

Beta Give feedback

No tasks being tracked yet.

@ashors1 has an example run showing that 2DP2TP4PP is now working, so we should re-run to get that baseline now (without TE first)

Then we should add TP=8 back when it's fixed

[terrykong]

Also, we should probably pytest.mark.skip some of the parametrizations for which we do not have baselines for so that when inspecting the CI output, we can know that the coverage is not complete:

JAX-Toolbox/.github/workflows/baselines/test_pax_mgmn_metrics.py

Line 31 in d577800

@pytest.mark.parametrize("baseline_filename", os.listdir(baselines_dir))