cifar workload error resolution

[init-22]

CIFAR workload error resolution: #889

[github-actions]

MLCommons CLA bot All contributors have signed the MLCommons CLA ✍️ ✅

[davidtweedle]

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

[init-22]

Oh thanks for pointing this out, I've updated the code which uses jax.tree.map now, can you please try it again?

[davidtweedle]

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) <#890 (comment)> 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 <#890 (comment)>, or unsubscribe <https://github.com/notifications/unsubscribe-auth/ACMTCD6H375WRDFT7MYINDL3UQLL3AVCNFSM6AAAAACHCQUF5WVHI2DSMVQWIX3LMV43OSLTON2WKQ3PNVWWK3TUHMZTGMZVGAZDSMBYHE> . You are receiving this because you commented.Message ID: ***@***.***>

[davidtweedle]

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) > <#890 (comment)> > > 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 > <#890 (comment)>, > or unsubscribe > <https://github.com/notifications/unsubscribe-auth/ACMTCD6H375WRDFT7MYINDL3UQLL3AVCNFSM6AAAAACHCQUF5WVHI2DSMVQWIX3LMV43OSLTON2WKQ3PNVWWK3TUHMZTGMZVGAZDSMBYHE> > . > You are receiving this because you commented.Message ID: > ***@***.***> >