ValueError: safe_zip() argument 2 is shorter than argument 1

[zwbx]

Thank you for your work. However, I encounter a bug when running the simple example.

The above exception was the direct cause of the following exception:
Traceback (most recent call last):
File "/home/wenbo/Documents/data/miniconda3/envs/octo/lib/python3.10/runpy.py", line 196, in _run_module_as_main
return _run_code(code, main_globals, None,
File "/home/wenbo/Documents/data/miniconda3/envs/octo/lib/python3.10/runpy.py", line 86, in _run_code
exec(code, run_globals)
File "/home/wenbo/.vscode/extensions/ms-python.debugpy-2024.6.0-linux-x64/bundled/libs/debugpy/adapter/../../debugpy/launcher/../../debugpy/main.py", line 39, in
cli.main()
File "/home/wenbo/.vscode/extensions/ms-python.debugpy-2024.6.0-linux-x64/bundled/libs/debugpy/adapter/../../debugpy/launcher/../../debugpy/../debugpy/server/cli.py", line 430, in main
run()
File "/home/wenbo/.vscode/extensions/ms-python.debugpy-2024.6.0-linux-x64/bundled/libs/debugpy/adapter/../../debugpy/launcher/../../debugpy/../debugpy/server/cli.py", line 284, in run_file
runpy.run_path(target, run_name="main")
File "/home/wenbo/.vscode/extensions/ms-python.debugpy-2024.6.0-linux-x64/bundled/libs/debugpy/_vendored/pydevd/_pydevd_bundle/pydevd_runpy.py", line 321, in run_path
return _run_module_code(code, init_globals, run_name,
File "/home/wenbo/.vscode/extensions/ms-python.debugpy-2024.6.0-linux-x64/bundled/libs/debugpy/_vendored/pydevd/_pydevd_bundle/pydevd_runpy.py", line 135, in _run_module_code
_run_code(code, mod_globals, init_globals,
File "/home/wenbo/.vscode/extensions/ms-python.debugpy-2024.6.0-linux-x64/bundled/libs/debugpy/_vendored/pydevd/_pydevd_bundle/pydevd_runpy.py", line 124, in _run_code
exec(code, run_globals)
File "/media/wenbo/12T/manipulation_project/lorax/examples/simple.py", line 36, in
lora_model(lora_params, jnp.ones((dim,)))
File "/home/wenbo/Documents/data/miniconda3/envs/octo/lib/python3.10/site-packages/qax/implicit/implicit_array.py", line 59, in implicit_f
outs_flat = f_wrapped.call_wrapped(*flat_args)
File "/home/wenbo/Documents/data/miniconda3/envs/octo/lib/python3.10/site-packages/jax/_src/linear_util.py", line 192, in call_wrapped
ans = self.f(*args, **dict(self.params, **kwargs))
File "/media/wenbo/12T/manipulation_project/lorax/examples/simple.py", line 10, in model
x = jax.nn.relu(x @ massive_w)
File "/home/wenbo/Documents/data/miniconda3/envs/octo/lib/python3.10/site-packages/jax/_src/numpy/array_methods.py", line 265, in deferring_binary_op
return binary_op(*args)
File "/home/wenbo/Documents/data/miniconda3/envs/octo/lib/python3.10/site-packages/qax/implicit/implicit_array.py", line 302, in process_primitive
outs = _default_handlers[primitive.name](primitive, *vals, params=params)
File "/home/wenbo/Documents/data/miniconda3/envs/octo/lib/python3.10/site-packages/qax/implicit/implicit_array.py", line 401, in _handle_pjit
outs = primitive.bind(*subfuns, *flat_inputs, **bind_params)
ValueError: safe_zip() argument 2 is shorter than argument 1

[zwbx]

I install all the dependent package. I also ran the pytest tests.py, it has the similar problem.

[zwbx]

Expect you reply if this project is still in maintenance

[davisyoshida]

Can you post the output of pip freeze?

[zwbx]

(lorax) project$ pip freeze
absl-py==2.1.0
beartype==0.18.5
build==1.2.1
CacheControl==0.14.0
certifi==2024.2.2
cffi==1.16.0
charset-normalizer==3.3.2
chex==0.1.86
cleo==2.1.0
crashtest==0.4.1
cryptography==42.0.7
distlib==0.3.8
dulwich==0.21.7
exceptiongroup==1.2.1
fastjsonschema==2.19.1
filelock==3.14.0
idna==3.7
importlib_metadata==7.1.0
iniconfig==2.0.0
installer==0.7.0
jaraco.classes==3.4.0
jax==0.4.28
jax-cuda12-pjrt==0.4.28
jax-cuda12-plugin==0.4.28
-e git+https://github.com/davisyoshida/lorax.git@1ebd0efa6c467974357bb9f73fe4a4967b487622#egg=jax_lorax
jaxlib==0.4.28
jeepney==0.8.0
keyring==24.3.1
markdown-it-py==3.0.0
mdurl==0.1.2
ml-dtypes==0.2.0
more-itertools==10.2.0
msgpack==1.0.8
numpy==1.24.4
nvidia-cublas-cu12==12.4.5.8
nvidia-cuda-cupti-cu12==12.4.127
nvidia-cuda-nvcc-cu12==12.4.131
nvidia-cuda-nvrtc-cu12==12.4.127
nvidia-cuda-runtime-cu12==12.4.127
nvidia-cudnn-cu12==8.9.7.29
nvidia-cufft-cu12==11.2.1.3
nvidia-cusolver-cu12==11.6.1.9
nvidia-cusparse-cu12==12.3.1.170
nvidia-nccl-cu12==2.21.5
nvidia-nvjitlink-cu12==12.4.127
opt-einsum==3.3.0
optax==0.1.9
packaging==24.0
pexpect==4.9.0
pkginfo==1.10.0
platformdirs==4.2.1
pluggy==1.5.0
plum-dispatch==2.3.6
poetry==1.8.3
poetry-core==1.9.0
poetry-plugin-export==1.8.0
ptyprocess==0.7.0
pycparser==2.22
Pygments==2.18.0
pyproject_hooks==1.1.0
pytest==7.4.4
qax==0.3.1
rapidfuzz==3.9.0
requests==2.31.0
requests-toolbelt==1.0.0
rich==13.7.1
scipy==1.9.3
SecretStorage==3.3.3
shellingham==1.5.4
tomli==2.0.1
tomlkit==0.12.5
toolz==0.12.1
trove-classifiers==2024.4.10
typing_extensions==4.11.0
urllib3==2.2.1
virtualenv==20.26.1
zipp==3.18.1

[zwbx]

poetry install has some incompatible problems, so I manually install some of them. I init the env by python 3.10

[davisyoshida]

@zwbx It looks like something broke with jax version 0.4.27. I'll see what needs to be changed. In the meanwhile, if you downgrade to jax==0.4.26, it should work.

[zwbx]

Thanks!