Issue 1:
The same code produces different outputs depending on whether it is run with a lax.while_loop or a native python while loop.

The output of the lax.while_loop is the incorrect one. We know this because this is part of a physics simulation tool (JAX-MD), and the lax.while_loop version of the code deletes some particles and duplicates others, while the python while loop version correctly preserves all particles.

This behavior is sensitive to the specific values in the multi-dimensional grid array that is being passed into the while loop. If grid is initialized randomly, the output of the two while loops becomes identical.

The lax.while_loop code does correctly preserve particles for other system configurations -- e.g. for a system which is 3d instead of 2d, or when the function is pmap-ed across multiple TPU cores.

We've verified this bug on both TPU and CPU.

This difference is output is present even when the while loop is run for only one step.

The code pasted below throws an assertion when the result of the two while loops are different. This repro is partially minimized, but we were unable to fully minimize it, since it seems brittle to specific aspects of the test case.

import jax.numpy as np

import jax
jax.config.update('jax_numpy_rank_promotion', 'raise')
from jax import lax, vmap

import math
import numpy as onp

import einops

# # Grid Packing / Unpacking

def get_grid_centers(box_size, cell_size, num_dims):
cells_count_per_side = int(onp.ceil(box_size / cell_size))

grid_centers = onp.zeros(tuple(
[cells_count_per_side]*num_dims + [num_dims]
))

for i in range(num_dims):
grid_centers[...,i] = (onp.reshape(onp.arange(cells_count_per_side),
tuple([1]*i + [cells_count_per_side] + [1]*(num_dims-i-1)))
+ 0.5)*cell_size

return np.array(grid_centers)

def positions_to_grid(position, box_size, cell_size, num_dims, particle_id, aux=None):
cells = int(math.ceil(box_size / cell_size))
N = position.shape[0]

cell_contents = np.concatenate([position, particle_id[:, None]], axis=-1)

if aux is not None:
cell_contents, *_ = set_aux(cell_contents, aux)

grid = np.zeros((cells ** num_dims, cell_contents.shape[-1]))
grid = grid.at[:N, :].set(cell_contents)

centers = get_grid_centers(box_size, cell_size, num_dims).reshape((-1, num_dims))
grid = grid.at[:N, :num_dims].add(-centers[:N])

grid = np.reshape(grid, (cells,) * num_dims + (-1,))

return grid

def pad_axis_channel_last(X, factors, axis_ind, max_grid_distance = 1):

num_dims = len(factors)
axes_s = ' '.join([chr(ord('x') + i) for i in range(num_dims)])
factors_s = ' '.join(['f'+chr(ord('x') + i) for i in range(num_dims)])
fact_d = {}
for i in range(num_dims):
fact_d['f' + chr(ord('x') + i)] = factors[i]

# all entries for all dimensions before the target axis
all_slice = tuple(slice(X.shape[j]) for j in range(axis_ind + 1))

# the end halo along the target axis
idx = all_slice + (slice(-max_grid_distance, X.shape[axis_ind + 1]),)
p = X[idx]
p = einops.rearrange(p, f'({factors_s}) {axes_s} c -> {factors_s} {axes_s} c', **fact_d)
p = np.roll(p, (0,) * axis_ind + (1,) + (0,) * (X.ndim - axis_ind - 1), axis=range(X.ndim))
p1 = einops.rearrange(p, f'{factors_s} {axes_s} c -> ({factors_s}) {axes_s} c')

# the start halo along the target axis
idx = all_slice + (slice(0, max_grid_distance),)
p = X[idx]
p = einops.rearrange(p, f'({factors_s}) {axes_s} c -> {factors_s} {axes_s} c', **fact_d)
p = np.roll(p, (0,) * axis_ind + (-1,) + (0,) * (X.ndim - axis_ind - 1), axis=range(X.ndim))
p2 = einops.rearrange(p, f'{factors_s} {axes_s} c -> ({factors_s}) {axes_s} c')

X = np.concatenate([p1, X, p2], axis=1 + axis_ind)

return X

def fold_factors(batch_size, grid_shape, max_grid_distance):
num_dims = len(grid_shape)
max_folds = onp.log(batch_size) / onp.log(2)

folds = onp.zeros((num_dims,), onp.int32)

for i in range(int(max_folds)*num_dims):
dim = i % num_dims
# folds[dim]+2, rather than +1, since we need an even number of cells for pairwise exchange in update_grid_locations
width, remain = divmod(grid_shape[dim], 2**(folds[dim]+2))
# width, remain = divmod(grid_shape[dim], 2**(folds[dim]+1))
# TODO is 2*max_grid_distance the right cutoff?
if remain == 0 and width >= max_grid_distance:
folds[dim] += 1
if onp.sum(folds) == max_folds:
break
factors = 2 ** folds

if onp.sum(folds) < max_folds:
print(f"Folds {folds} have sum smaller than target of {max_folds}. This corresponds to a batch size of {onp.prod(factors)} rather than the target of {batch_size}.")

return factors

def order_grid_by_factors(num_dims):
front_order = ()
back_order = ()

for i in range(num_dims):
front_order += (2 * i,)
back_order += (2 * i + 1,)

return front_order + back_order + (2 * num_dims,)

def fold_grid(X, max_grid_distance, batch_size=None, factors=None, inner_fold=True):
num_dims = X.ndim - 1

if factors is None:
factors = fold_factors(batch_size, X.shape[:-1], max_grid_distance)

X_shape = []
for i in range(num_dims):
_, ragged = divmod(X.shape[i], factors[i])
assert ragged == 0
X_shape += [factors[i], X.shape[i] // factors[i]]
X_shape += [X.shape[-1]]

if onp.min(onp.array(X.shape[:-1]) / onp.array(factors)) < max_grid_distance:
print(f"Folded grid may be too small. This failure is happening when splitting across devices. factors {factors} X.shape {X.shape} new X_shape {X_shape}.")

ordering = order_grid_by_factors(num_dims)

X = np.reshape(X, X_shape)
X = np.transpose(X, ordering)

if inner_fold:
X = X.reshape((-1,) + X.shape[num_dims:])
else:
X = folded_pad(X, max_grid_distance, factors)

return X, factors

# Dynamics

def pairwise_exchange(X, cell_size, axis, num_dims):
particle_data_size = X.shape[-1] - num_dims

axes_pre = ' '.join(['(' + chr(ord('x') + i) + ' q)' if i == axis else chr(ord('x') + i) for i in range(num_dims)])
axes_post = ' '.join([chr(ord('x') + i) for i in range(num_dims)])
X = einops.rearrange(X, f'{axes_pre} c -> q {axes_post} c', q=2)

shift = onp.zeros((2, num_dims+particle_data_size))
shift[0,axis] = cell_size
shift[1,axis] = -cell_size
shift = shift.reshape((2,) + (1,)*num_dims + (num_dims+particle_data_size,))
_Xrev = X[::-1]
X_swapped = _Xrev + shift*(_Xrev[..., [-1]]>0) # only add offset to occupied cells

square_displacement = lambda _X: np.sum(np.sum(_X[...,:num_dims]**2 * (_X[...,[-1]]>0), axis=-1, keepdims=True), axis=0, keepdims=True)

keep_mask = (square_displacement(X) < square_displacement(X_swapped))

X_new = X*keep_mask + X_swapped*(~keep_mask)

X_new = einops.rearrange(X_new, f'q {axes_post} c -> {axes_pre} c')

return X_new

pairwise_exchange = vmap(pairwise_exchange, in_axes=(0, None, None, None))

def update_grid_locations(X, cell_size, num_dims, factors):
"""
Move particles to new grid cells if any coordinate entries in X are larger
magnitude than cell_size/2 (ie, if particle has left current grid cell)
(this could also be phrased as a conv, but I expect it would be less efficient + more confusing)
"""

axis = 1
X = pad_axis_channel_last(X, factors, axis)
X = pairwise_exchange(X, cell_size, axis, num_dims)
body_slice = tuple(slice(X.shape[j]) for j in range(axis+1)) + (slice(1,-1),)
X = X[body_slice]

return X, None

def test_grid(grid_topology=(), num_dims=2):

cell_size = 1./4.
max_grid_distance = 5
box_size = 512.

points = onp.linspace(0., 50, num=66)
X, Y = onp.meshgrid(points, points)
R = onp.stack((X.ravel(), Y.ravel()), axis=1) + 0.1

R = onp.array(R, onp.float64)

grid = positions_to_grid(R, box_size, cell_size, num_dims, np.arange(R.shape[0]) + 1)

grid, factors = fold_grid(grid, max_grid_distance,
batch_size=128)

def cond_fn(grid_old_grid):
grid, i = grid_old_grid
return (i < 1)

def body_fn(grid_old_grid):
grid, i = grid_old_grid
_grid, overflow_count = update_grid_locations(grid, cell_size, num_dims, factors)
return _grid, i+1

grid_old_grid = (grid, 0)

### BUG IS HERE. Compute a while loop two different ways. Result should be the
### same. Instead it's different.

new_tuple_lax = lax.while_loop(cond_fn, body_fn, grid_old_grid)

new_tuple_while = grid_old_grid
while cond_fn(new_tuple_while):
new_tuple_while = body_fn(new_tuple_while)

max_dev = onp.max(onp.abs(new_tuple_lax[0] - new_tuple_while[0]))
print(max_dev)
assert max_dev == 0

test_grid()

Issue 2:
Hi all,

I was discussing with @jakevdp in #6466 about adding support for sparse matrices in Jax. He already did quite some groundwork, but there's a lot of open questions on the what and how. Quoting Jake here:

2D CSC/CSR/COO may not be the abstraction we want in the end... pydata-sparse has N-dimensional sparse arrays implemented using COO or a multi-dimensional generalization of CSR. Similarly taco has routines for general N-dimensional sparse computation. That route may be better, particularly for some deep learning workflows that e.g. have 4D batched arrays of sparse image sequences.

XLA scatter/gather implementations only go so far... any operations we implement should have efficient lowerings on as many backends as possible - I've started with cusparse on GPU because the low-level routines are already available in our GPU build process. This means that every non-trivial operation should be either implemented as a primitive or composed of other primitives

JAX is about composability, and in particular pmap/vmap, jit, and grad are central to the API. I want to spend some time exploring how we can best implement gradients of sparse operations or batched sparse operations, and also think about how ongoing work on dynamic shapes can be utilized to make jit of sparse operations more efficient and usable. With all that up in the air I want to keep the design minimal for now in case we need to re-think our approach based on what we find

all of the above should be motivated by real-world workflows that people have. I'm starting to gather a few example workflows so we can have that in mind.

So I set out to compare the implementations in pytorch and tensorflow to get a feel for what's out there, and considering the points above. Specifically, we're interested in how they implement it low-level: do they secretly turn everything dense under the surface, or are they given the full treatment with specific ops?

I'm summarizing what I found in the table below - TL;DR Pytorch offers better support. Since documentation and support are all over the place, in this notebook you can find a bunch of common operations with sparse tensor to see if they are supported. It's been a while since I did tensorflow, so let me know if I made a mistake!

Feature
Pytorch
Tensorflow

Documentation
here
here

Formats
COO
COO

Supported ops
variations of matmul
sums, min, matmul, element-wise ops through tf.sparse.map_values

Specialized ops
Probably1
Seems so 2

Grad of sparse3
Limited to specific matmul
No.

Sparse grad4
Yes
No.

Dimensions
2D, Hybrid
2D - tensor supports higher-D, but operations not

Uncoalesced5allowed
Yes
No mention

Sparse - sparse matmul
No
No

Extra goodies
Has an adam optimizer for sparse tensor
Nope

To do

Dig in the source code a little further (although if someone who actually is familiar with the code could chime in, that'd be great.)

Check how pydata-sparse implement n-D sparse arrays.

Proposal:

Get a very minimal differentiable sparse matmul working for COO format using XLA ops and one GPU-specific backend in cusparse (some help and guidance here would be nice)

I figured we can use this issue to track progress and discuss, so ... thoughts?

Footnotes

They mention Torchcudasparse is now defunct here but I can't find any other mention in the source code. ↩

Sparse ops are imported from gen_sparse_ops.py, which doesnt exist in the repo? ↩

Meaning gradient w.r.t to the sparse tensor ↩

Meaning calculating the gradient without turning it into a dense tensor ↩

Uncoalesced means duplicate entries in the coordinates - total value of the element is then sum of duplicate elements. ↩