Update forward_pass_logit_checker for new dimension

[A9isha]

Description

There is an added dimension for jax.experimental.multihost_utils.process_allgather(full_train_logits), so updating forward_pass_logit_checker for new dimension

FIXES: b/389749060

Tests

Locally reran forward_pass_logit_checker for mistral

Checklist

Before submitting this PR, please make sure (put X in square brackets):

• I have performed a self-review of my code.
• I have necessary comments in my code, particularly in hard-to-understand areas.
• I have run end-to-end tests tests and provided workload links above if applicable.
• I have made or will make corresponding changes to the doc if needed.

[gobbleturk]

Awesome thanks Anisha!

Don't we have to check for jax version, or does this change support both old and new? Will the tests targeting jax stable break with this change?
E.g. in the past sometimes we have had code like

if jax.version >= X
    do_new_thing
else:
    do_old_thing

Is that needed here?

[A9isha]

Thanks Matt. I thought further comparing the tests for Mistral and Llama2/Gemma2 and realized Jax actually handles dimension of value 1, so this updated way is a cleaner solution

Awesome thanks Anisha!

Don't we have to check for jax version, or does this change support both old and new? Will the tests targeting jax stable break with this change? E.g. in the past sometimes we have had code like

if jax.version >= X
    do_new_thing
else:
    do_old_thing

Is that needed here?

[A9isha]

Jax actually handles dimension of value 1, so this updated way is a cleaner solution

@AlipMagical Could you please give a bit more context?