To install specific CUDA version builds for PyTorch, use
pip uninstall -y torch pip install --no-cache-dir torch -f https://download.pytorch.org/whl/cu101/torch_stable.html
cu101 corresponds to CUDA 10.1, and needs to be changed appropriately. Test for GPU
import torch print(torch.cuda.is_available())
Similarly, for JAX:
pip uninstall -y jaxlib pip install --no-cache-dir jaxlib==0.1.64+cuda111 -f https://storage.googleapis.com/jax-releases/jax_releases.html
Test for GPU with,
```python from jax.lib import xla_bridge print(xla_bridge.get_backend().platform)
For TensorFlow, verify the CUDA-packaged wheels here and install the appropriate version.
Test for GPU devices with,
import tensorflow as tf print(tf.config.list_physical_devices('XLA_GPU')) ## or 'GPU'
If widgets like tqdm progress bars are not working in Jupyter Lab, run the following once in every new environment containing Jupyter.
$ pip install ipywidgets $ jupyter nbextension enable --py widgetsnbextension $ jupyter labextension install @jupyter-widgets/jupyterlab-manager
After searching innumerable times, I finally have it handy. This was originally discussed in Issue #394 on Github.
For , the pairwise distance matrix between each pair of these batched matrices is , where ... represent arbitrary batch dimension (think batches of pairs of and samples of dimension ).
## in PyTorch def pairwise_dist(x, y): xx = (x * x).sum(dim=-1).unsqueeze(-1) yy = (y * y).sum(dim=-1).unsqueeze(-2) xy = torch.einsum('...ji,...ki->...jk', x, y) d = xx + yy - 2. * xy return d