Linting low-hanging fruit across MaxText, MaxText/input_pipeline, MaxText/layers

[SamuelMarks]

Description

For commentary see #1108

EDIT0: Hmm this has 194 files… can think about how to split it out smaller, maybe per folder rather than per module?

EDIT1: Here we go again

fd -Ftd -j1 -x bash -c 'd=${0:2}; d=${d//\//.}; git switch main && 
if ! git rev-parse --verify &>/dev/null; then git checkout -b "PKG-FIX_${d}" && git checkout test-fixer "${0}" && git commit -S -am "[${d}] Python package layout support and minor other fixes" && git push mine "PKG-FIX_${d}" && gh pr create --title "[${d}] Python package layout support and minor other fixes" --body "
For commentary see #1108"; fi' {//}

EDIT2: Merge this before #1108 then #1108 might be small-enough to merge. Just tell me and I'll squash and remove any reference to constants.py.

Tests

N/A

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.

[SamuelMarks]

@shralex The new argument so that inference_metadata can make it all the way into main otherwise it's lost (and that's a bug I assume?)

[bvandermoon]

Thank you @SamuelMarks. Generally LGTM, just a couple comments/questions

[bvandermoon]

Why is the location of prefix_cache changing from jetstream.core to MaxText? Also a minor nit: If you are going to inline some of the MaxText imports like max_utils, prefix_cache, it might be good to also inline maxengine and pyconfig

[SamuelMarks]

@bvandermoon Oh my earlier code-base wide find/replace looks at the file hierarchy as the module hierarchy so MaxText/prefix_cache.py means that import prefix_cache became from MaxText import prefix_cache. I can change this one to jetstream.core. - What about MaxText/inference_microbenchmark.py and MaxText/tests/prefix_cache_test.py? - I assume they stay as is.

EDIT: #1643 asked more broadly