Enhancement: Maybe try using Grain to load data for JAX

[imgengineer]

I'm glad to see another 3DGS project written in JAX, which saves me the pain of compiling CUDA. Regarding data input, I'd like to offer a small suggestion: you can try using grains to load data.

[takah29]

Thank you for your comment. I'm glad you're interested in the project.

Actually, I did use Grain in the early stages, but I replaced it with my own implementation since the loading was straightforward and didn't require any preprocessing.

Looking at it again, the code would probably be cleaner if implemented with Grain. I'll try rewriting it when I have time.