Conflicting dependencies between Jax and neural-tangents

[snehjp2]

Hi, I am running Jax v0.4.1 (for CPU) and neural-tangents v0.5.0. When importing neural-tangents in a script, I encounter this error:

ImportError: cannot import name 'tree_multimap' from 'jax.tree_util'

After some googling, I thought it may have to do with the version of neural-tangents I was running. So I attempted to install the 'bleeding edge' version from the GitHub using:

git clone https://github.com/google/neural-tangents; cd neural-tangents pip install -e .

The install was not successful and resulted in the following error message:

ERROR: Cannot install neural-tangents because these package versions have conflicting dependencies. ERROR: ResolutionImpossible: for help visit https://pip.pypa.io/en/latest/topics/dependency-resolution/#dealing-with-dependency-conflicts

I'm not sure how to proceed here; any tips? I am using a 2021 M1 Mac.

[romanngg]

EDIT: just read your second part - will investigate.

Old reply: Could you try updating NT to v0.6.1? 0.5.0 is pretty old, so may be incompatible with newer JAX releases.

[snehjp2]

When I assert the version of NT to v0.6.1, it yields the following conflicting dependencies:

The conflict is caused by: tf2jax 0.3.1 depends on tensorflow>=2.8.0 tf2jax 0.3.0 depends on tensorflow>=2.8.0

When trying to install tensor flow using python -m pip install tensorflow-macos (I do not use tensorflow so did not have it installed), I get a error: legacy-install-failure. 🙃

I had used NT / Jax several months ago and did not have any problems / did not need to have tensorflow installed. Is there a workaround to that?

[romanngg]

Looks like there's an issue with TF on M1 Macs. NT replies on TF only for the nt.experimental module; if you don't use it, a quick workaround could be to comment out these two lines on your machine locally (after cloning the repo):

https://github.com/google/neural-tangents/blob/db1c240af9f04453d2470b173af4a4717403f90a/setup.py#L32

and

https://github.com/google/neural-tangents/blob/db1c240af9f04453d2470b173af4a4717403f90a/neural_tangents/__init__.py#L21

Then you can run pip install -e . and I think NT would work OK (but lack the experimental module). I'll think how to handle this more gracefully.

[snehjp2]

Hi,

So the simple git clone of the repo, using git clone https://github.com/google/neural-tangents keeps defaulting to v0.5.0 as shown in __init__.py. I don't believe there is a way to assert the version to install unless I'm using a pip install, but when I did that it was breaking because of the tensorflow dependency.

[romanngg]

Hm can this be an issue with a previous install? Could you try

pip uninstall neural-tangents
pip install git+https://github.com/google/neural-tangents.git@no_experimental

(I've created the branch with experimental deps commented out) The versions should be displayed as 0.6.2

[romanngg]

Should be fixed as of v0.6.2!