Merge main (#168)

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

* Add compilation metrics, script to assert metrics, add conditional import of hugginface_hub and update huggingface hub

Signed-off-by: Kunjan patel <kunjanp@google.com>

* Formatting

Signed-off-by: Kunjan patel <kunjanp@google.com>

* Fix pylint issues

Signed-off-by: Kunjan patel <kunjanp@google.com>

---------

Signed-off-by: Kunjan patel <kunjanp@google.com>

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

* Fix 413751357 updating orbax version which is lowering the JAX version

* Updating the right dependency

---------

Signed-off-by: Kunjan patel <kunjanp@google.com>
Co-authored-by: Param Bole <parambole@google.com>

Author: coolkp

requirements.txt

@@ -19,7 +19,7 @@ Pillow
 pylint
 pyink
 pytest==8.2.2
-tensorflow==2.17.0
+tensorflow>=2.17.0
 tensorflow-datasets>=4.9.6
 ruff>=0.1.5,<=0.2
 git+https://github.com/mlperf/logging.git