Numerical code snippets

🚀 tech
Table of Contents

CUDA Specific Builds


To install specific CUDA version builds for PyTorch, use

pip uninstall -y torch
pip install --no-cache-dir torch -f

cu101 corresponds to CUDA 10.1, and needs to be changed appropriately. Test for GPU with

import torch


Similarly, for JAX:

pip uninstall -y jaxlib
pip install --no-cache-dir jaxlib==0.1.64+cuda111 -f

Test for GPU with,

from jax.lib import xla_bridge


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'

Jupyter Lab widgets

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.

Vectorized Pairwise Distances

For XR...×m×d,YR...×n×d\mathbf{X} \in \mathbb{R}^{... \times m \times d}, \mathbf{Y} \in \mathbb{R}^{... \times n \times d}, the pairwise distance matrix between each pair of these batched matrices is DR...×m×n\mathbf{D} \in \mathbb{R}^{... \times m \times n}, where … represent arbitrary batch dimension (think batches of pairs of mm and nn samples of dimension dd).

## 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,>...jk', x, y)
    d = xx + yy - 2. * xy
    return d