Can’t import both jax and tensorflow (causes kernel restart)

[zeevikal]

hi,
after importing those libs:

import pandas as pd
import jax.numpy as np
import tensorflow as tf

I tried to run this cell and the results:

[mattjj]

Currently we can’t import both jax and tensorflow. I’m not sure the exact reason but it has something to do with how we build XLA. Definitely a bug we need to fix, but for now there’s no workaround that I know of.

[mattjj]

By the way, unless you have additional info to suggest otherwise, I think this is not an OOM. I’ll revise the issue title to reflect my best understanding (that this is about importing both jax and tf). Let me know if we need to change it back.

Thanks for bringing this up!

[zeevikal]

thanks for quick response @mattjj .
I tried to run the example in jax repo with adding tf lib and got the following error:

I think the bug is related to C++ code.
maybe a program where it calls free() on a pointer that wasn't obtained with malloc(). or maybe some other memory-related bug in the code that triggers this behavior.

[mattjj]

I think the issue stems from the build process, rather than a bug in the code itself. We build XLA as part of tensorflow, but that means our XLA "jaxlib" .so file and tensorflow's .so files may have some symbol redundancies, or something.

@hawkinsp any thoughts to jot down here?

[hawkinsp]

Just to follow up: I finally had time to look at this. It turns out that the jaxlib (i.e. XLA) Python extension depends on a few small pieces of TensorFlow (mostly some C++ support libraries). If both the jaxlib and tensorflow Python extensions export the same symbols, a crash occurs. The fix is to restrict the set of symbols we export from the XLA python extension; there's no good reason for the XLA python extension to also export bits of TF.

I'll push a fix shortly.

[mattjj]

FYI @rsepassi

[hawkinsp]

PR #265 checked in a fix to jaxlib that should fix this problem.

Note that for the moment you'll either need to rebuild jaxlib from source, or wait until we build and push new Jaxlib wheels.

Hope that helps!

[rsepassi]

awesome!

…

On Fri, Jan 18, 2019 at 12:01 PM Peter Hawkins ***@***.***> wrote: PR #265 <#265> checked in a fix to jaxlib that should fix this problem. Note that for the moment you'll either need to rebuild jaxlib from source, or wait until we build and push new Jaxlib wheels. Hope that helps! — You are receiving this because you were mentioned. Reply to this email directly, view it on GitHub <#120 (comment)>, or mute the thread <https://github.com/notifications/unsubscribe-auth/ABEGW4Ved-Oh2CyvWk4rRKk7aAnlVmwbks5vEigygaJpZM4ZVFtK> .

[hawkinsp]

I've built and pushed updated jaxlib wheels (0.1.6) to pypi. If you update your jaxlib version (pip install --upgrade jaxlib, if you aren't using a CUDA build), you should be able to import both tensorflow and jax with no problems.