apax might be incompatible with latest jax release

[PythonFZ]

With the latest JAX I get AttributeError: 'EvalTrace' object has no attribute 'level'

in

│ /data/fzills/tools/apax/apax/train/checkpoints.py:57 in create_params                                                                                                       │
│                                                                                                                                                                             │
│    54 │   log.info(f"initializing {n_models} model(s)")                                                                                                                     │
│    55 │                                                                                                                                                                     │
│    56 │   if n_models == 1:                                                                                                                                                 │
│ ❱  57 │   │   params = model.init(model_rng[0], *sample_input)                                                                                                              │
│    58 │   elif n_models > 1:                                                                                                                                                │
│    59 │   │   num_args = len(sample_input)                                                                                                                                  │
│    60 │   │   # vmap only over parameters, not over any data from the input

[M-R-Schaefer]

I'll look into it next week

[M-R-Schaefer]

closed by #391