Merge main#168

Merged

coolkp merged 2 commits intokunjan-devfrom

Apr 28, 2025

Commits on Apr 25, 2025

Add compilation metrics, parse step time in e2e, add name to train partial etc (#166)
coolkp
authored

Commits on Apr 26, 2025

Fix 413751357 updating version of depdendency which is lowering the JAX version (#167)
parambole
authored