Fix MaxText pyconfig import in page_manager_test

[bvandermoon]

Description

Saw this error while running python3 -m pytest on a local TPU VM:

ImportError while importing test module '/home/bvandermoon/workspace/nnx_maxtext/maxtext/MaxText/tests/inference/page_manager_test.py'.
Hint: make sure your test modules/packages have valid Python names.
Traceback:
MaxText/tests/inference/page_manager_test.py:21: in <module>
    import pyconfig
E   ModuleNotFoundError: No module named 'pyconfig'

Updating this import let the tests start locally

Tests

After this change, I was able to run the unit tests on a TPU VM with this:

python3 -m pytest
python3 -m pytest MaxText/tests/inference/page_manager_test.py # Tests passed

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.

[khatwanimohit]

I think there's a confusion here. The import pyconfig on line 21 of MaxText/tests/inference/paged_attention_test.py acutally is not the pypi package. It is MaxText/pyconfig.py. The import in that test should be made from MaxText import pyconfig.

[bvandermoon]

I think there's a confusion here. The import pyconfig on line 21 of MaxText/tests/inference/paged_attention_test.py acutally is not the pypi package. It is MaxText/pyconfig.py. The import in that test should be made from MaxText import pyconfig.

Updated, thank you @khatwanimohit

[patemotter]

Thanks for catching/fixing this.