Reduce numerical difference in llama4 vision encoder test

[aireenmei]

Description

• Update encoder layer number from 6 to 34, matching llama4 vision encoder
• Update the allclose criteria rtol from 1e-3 to 0.01, keep atol the same 0.05. Since rtol=1e-3, atol=0.05 is usually used for single layer testing, we loose this a bit for testing the 34-layer vision encoder. The end to end logit_checker uses rto=0.1, atol=0.1
• Update precision and make the dummy test input in a more realistic range
  • before this change:

E           AssertionError: 
E           Not equal to tolerance rtol=0.01, atol=0.05
E           
E           Mismatched elements: 306285 / 3249664 (9.43%)
E           Max absolute difference among violations: 0.2326374
E           Max relative difference among violations: 3.343927
E            ACTUAL: array([[[11.217089, 12.195944,  6.685423, ...,  8.819061, 13.093238,
E                     8.150303],
E                   [11.21041 , 12.189107,  6.688676, ...,  8.817145, 13.097935,...
E            DESIRED: array([[[11.303663, 12.14169 ,  6.760583, ...,  8.844138, 13.108282,
E                     8.161341],
E                   [11.303666, 12.141687,  6.760585, ...,  8.844134, 13.108283,...

• After this change: pass

If the change fixes a bug or a Github issue, please include a link, e.g.,:
FIXES: b/420952122

Notice 1: Once all tests pass, the "pull ready" label will automatically be assigned.
This label is used for administrative purposes. Please do not add it manually.

Notice 2: For external contributions, our settings currently require an approval from a MaxText maintainer to trigger CI tests.

Tests

Please describe how you tested this change, and include any instructions and/or
commands to reproduce.

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.

[ZhaoyueCheng]

Thanks Aireen for looking into this issue!

[hengtaoguo]

Great changes, thanks!