cifar workload error resolution
CIFAR workload error resolution: #889
MLCommons CLA bot All contributors have signed the MLCommons CLA ✍️ ✅
Hi Isaac, Thank you for doing this. When I tried it in my docker container it said
File "/algorithmic-efficiency/algoperf/workloads/cifar/cifar_jax/workload.py", line 223, in _normalize_eval_metrics
return jax.tree_map(lambda x: x / num_examples, total_metrics)
^^^^^^^^^^^^
File "/usr/local/lib/python3.11/site-packages/jax/_src/deprecations.py", line 54, in getattr
raise AttributeError(message)
AttributeError: jax.tree_map was removed in JAX v0.6.0: use jax.tree.map (jax v0.4.25 or newer) or jax.tree_util.tree_map (any JAX version).
When I change this line in workload.py to jax.tree.map(....) everything works fine. If it is just that I am using the wrong jax version, then you can just ignore this problem. Anyways, thanks again - David
Oh thanks for pointing this out, I've updated the code which uses jax.tree.map now, can you please try it again?
OK I will try it again
On Thu, Sep 25, 2025 at 12:41 PM Isaac @.***> wrote:
init-22 left a comment (mlcommons/algorithmic-efficiency#890) https://github.com/mlcommons/algorithmic-efficiency/pull/890#issuecomment-3335029089
Oh thanks for pointing this out, I've updated now which uses jax.tree.map, can you please try it again now?
— Reply to this email directly, view it on GitHub https://github.com/mlcommons/algorithmic-efficiency/pull/890#issuecomment-3335029089, or unsubscribe https://github.com/notifications/unsubscribe-auth/ACMTCD6H375WRDFT7MYINDL3UQLL3AVCNFSM6AAAAACHCQUF5WVHI2DSMVQWIX3LMV43OSLTON2WKQ3PNVWWK3TUHMZTGMZVGAZDSMBYHE . You are receiving this because you commented.Message ID: @.***>
I tried this new code with my jax version (you are right I am using the jax from pyproject.toml) and everything is working now. Thanks!
On Thu, Sep 25, 2025 at 12:54 PM David Tweedle @.***> wrote:
OK I will try it again
On Thu, Sep 25, 2025 at 12:41 PM Isaac @.***> wrote:
init-22 left a comment (mlcommons/algorithmic-efficiency#890) https://github.com/mlcommons/algorithmic-efficiency/pull/890#issuecomment-3335029089
Oh thanks for pointing this out, I've updated now which uses jax.tree.map, can you please try it again now?
— Reply to this email directly, view it on GitHub https://github.com/mlcommons/algorithmic-efficiency/pull/890#issuecomment-3335029089, or unsubscribe https://github.com/notifications/unsubscribe-auth/ACMTCD6H375WRDFT7MYINDL3UQLL3AVCNFSM6AAAAACHCQUF5WVHI2DSMVQWIX3LMV43OSLTON2WKQ3PNVWWK3TUHMZTGMZVGAZDSMBYHE . You are receiving this because you commented.Message ID: @.***>