Migrate ArrayRecord parsing off TF

[aireenmei]

Description

The ParseFeatures function used by Grain pipeline use tf.io.parse_example to parse examples, which uses the proto from the tensorflow library. To remove tensorflow dependency, I download the proto files from tensorflow library, modify the path and compile. The proto files and the compiled *_pb.py files are now in input_pipeline/protos. Migrate ParseFeatures to use the input_pipeline/protos instead of tf. Also in input_pipeline_utils.py, move tensorflow from top level import to local import in the corresponding functions.

Steps:

1. example.proto and feature.proto are downloaded from https://github.com/tensorflow/tensorflow/tree/master/tensorflow/core/example
2. Compile with cd /mnt/disks/pd2/maxtext/src && protoc --proto_path=. --python_out=. maxtext/input_pipeline/protos/feature.proto maxtext/input_pipeline/protos/example.proto, this generates example_pb2.py and feature_pb2.py
3. Update input_pipeline_utils.py to use protos/example_pb2

Tests

• CI test includes testing grain pipeline reading arrayrecord files
• Manually test to compare step time, found no change. On v4-8, test this script
  Before PR:

INFO:absl:completed step: 0, seconds: 19.222, TFLOP/s/device: 8.482, Tokens/s/device: 1278.556, total_weights: 97214, loss: 10.877
INFO:absl:To see full metrics 'tensorboard --logdir=gs://aireenmei-multipod/test/grain/grain-2026-03-04-02-56-46/tensorboard/'
INFO:absl:completed step: 1, seconds: 0.059, TFLOP/s/device: 2771.846, Tokens/s/device: 417802.863, total_weights: 96940, loss: 10.864
INFO:absl:completed step: 2, seconds: 1.385, TFLOP/s/device: 117.687, Tokens/s/device: 17739.127, total_weights: 96752, loss: 9.764
INFO:absl:completed step: 3, seconds: 0.972, TFLOP/s/device: 167.659, Tokens/s/device: 25271.315, total_weights: 97281, loss: 9.248
INFO:absl:completed step: 4, seconds: 1.202, TFLOP/s/device: 135.639, Tokens/s/device: 20444.954, total_weights: 97461, loss: 9.082
INFO:absl:completed step: 5, seconds: 1.203, TFLOP/s/device: 135.493, Tokens/s/device: 20423.054, total_weights: 95557, loss: 8.862
INFO:absl:completed step: 6, seconds: 1.210, TFLOP/s/device: 134.745, Tokens/s/device: 20310.274, total_weights: 97460, loss: 8.737
INFO:absl:completed step: 7, seconds: 1.195, TFLOP/s/device: 136.425, Tokens/s/device: 20563.522, total_weights: 97706, loss: 8.536
INFO:absl:completed step: 8, seconds: 1.202, TFLOP/s/device: 135.675, Tokens/s/device: 20450.347, total_weights: 96937, loss: 8.510
INFO:absl:completed step: 9, seconds: 1.202, TFLOP/s/device: 135.667, Tokens/s/device: 20449.224, total_weights: 97317, loss: 8.446

After PR

INFO:absl:completed step: 0, seconds: 16.573, TFLOP/s/device: 9.838, Tokens/s/device: 1482.864, total_weights: 97214, loss: 10.877
INFO:absl:To see full metrics 'tensorboard --logdir=gs://aireenmei-multipod/test/grain/grain-2026-03-04-02-59-25/tensorboard/'
INFO:absl:completed step: 1, seconds: 0.056, TFLOP/s/device: 2920.236, Tokens/s/device: 440169.792, total_weights: 96940, loss: 10.864
INFO:absl:completed step: 2, seconds: 1.391, TFLOP/s/device: 117.180, Tokens/s/device: 17662.583, total_weights: 96752, loss: 9.764
INFO:absl:completed step: 3, seconds: 0.969, TFLOP/s/device: 168.206, Tokens/s/device: 25353.909, total_weights: 97281, loss: 9.248
INFO:absl:completed step: 4, seconds: 1.201, TFLOP/s/device: 135.703, Tokens/s/device: 20454.568, total_weights: 97461, loss: 9.082
INFO:absl:completed step: 5, seconds: 1.203, TFLOP/s/device: 135.538, Tokens/s/device: 20429.760, total_weights: 95557, loss: 8.862
INFO:absl:completed step: 6, seconds: 1.204, TFLOP/s/device: 135.469, Tokens/s/device: 20419.355, total_weights: 97460, loss: 8.737
INFO:absl:completed step: 7, seconds: 1.201, TFLOP/s/device: 135.751, Tokens/s/device: 20461.806, total_weights: 97706, loss: 8.536
INFO:absl:completed step: 8, seconds: 1.202, TFLOP/s/device: 135.601, Tokens/s/device: 20439.207, total_weights: 96937, loss: 8.510
INFO:absl:completed step: 9, seconds: 1.202, TFLOP/s/device: 135.595, Tokens/s/device: 20438.408, total_weights: 97317, loss: 8.446

Checklist

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

• I have performed a self-review of my code. For an optional AI review, add the gemini-review label.
• 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, including adding new documentation pages to the relevant Table of Contents (toctree directive) as explained in our documentation.

[codecov]

Codecov Report

❌ Patch coverage is 45.45455% with 42 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
src/maxtext/input_pipeline/protos/feature_pb2.py	27.50%	28 Missing and 1 partial ⚠️
src/maxtext/input_pipeline/protos/example_pb2.py	63.15%	6 Missing and 1 partial ⚠️
src/maxtext/input_pipeline/input_pipeline_utils.py	66.66%	4 Missing and 2 partials ⚠️

📢 Thoughts on this report? Let us know!

[gagika]

a few minor comments