Matthew Johnson
Matthew Johnson
Thanks so much for digging in! Kayhan, good to hear from you! Looks like you’ve already found this code is currently only Python 2.7 compatible. The fmap stuff is vestigial...
Is this consistent with our "remove x64 flag" plans? Or e.g. should we just switch the defaults to be 32bit?
Hrm, I don't know why this would happen, though I would guess it's an effect of backend-initialization. That is, when you run the first JAX operation (_not_ when JAX is...
> Ideally I guess one wouldn't expect gpu backend be initialized when working only on the cpu backend. We might be able to make this initialization lazier... @skye do you...
@superbobry any thoughts or references to point to? (@billmark is a Googler.)
https://github.com/google/jax/pull/9118 may have added this context manager already, but I'm not sure. Checking with @skye ...
To put a fine point on it, the analogue of code like this: ```python with torch.cuda.device(2): # specify default device ... # create array on the default device ``` Is...
I think that is *not* intended. (@ayaka14732 thanks for following up and for the pings. Sorry I didn't notice your replies until now!)
Sorry for prematurely closing this. I really should've added a test, and then I would've seen that it doesn't work! On CPU at least, it looks like `xla_extension.Client.buffer_from_pyval(x, None)` doesn't...
https://github.com/google/jax/pull/12048 will add an API for clearing caches.