<ul dir="auto">
<li>The same code produces different outputs depending on whether it is run with a <code class="notranslate">lax.while_loop</code> or a native python <code class="notranslate">while</code> loop.</li>
<li>The output of the <code class="notranslate">lax.while_loop</code> is the incorrect one. We know this because this is part of a physics simulation tool (JAX-MD), and the <code class="notranslate">lax.while_loop</code> version of the code deletes some particles and duplicates others, while the python <code class="notranslate">while</code> loop version correctly preserves all particles.</li>
<li>This behavior is sensitive to the specific values in the multi-dimensional <code class="notranslate">grid</code> array that is being passed into the while loop. If <code class="notranslate">grid</code> is initialized randomly, the output of the two while loops becomes identical.</li>
<li>The <code class="notranslate">lax.while_loop</code> 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.</li>
<li>We've verified this bug on both TPU and CPU.</li>
<li>This difference is output is present even when the while loop is run for only one step.</li>
</ul>
<p dir="auto">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.</p>
<div class="highlight highlight-source-python notranslate position-relative overflow-auto" dir="auto" data-snippet-clipboard-copy-content="
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 -&gt; {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 -&gt; ({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 -&gt; {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 -&gt; ({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 &gt;= max_grid_distance:
      folds[dim] += 1
    if onp.sum(folds) == max_folds:
      break
  factors = 2 ** folds

  if onp.sum(folds) &lt; max_folds:
    print(f&quot;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}.&quot;)

  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)) &lt; max_grid_distance:
    print(f&quot;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}.&quot;)

  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 -&gt; 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]]&gt;0) # only add offset to occupied cells

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

  keep_mask = (square_displacement(X) &lt; square_displacement(X_swapped))

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

  X_new = einops.rearrange(X_new, f'q {axes_post} c -&gt; {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):
  &quot;&quot;&quot;
  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)
  &quot;&quot;&quot;

  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 &lt; 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()"><pre class="notranslate"><span class="pl-k">import</span> <span class="pl-s1">jax</span>.<span class="pl-s1">numpy</span> <span class="pl-k">as</span> <span class="pl-s1">np</span>

<span class="pl-k">import</span> <span class="pl-s1">jax</span>
<span class="pl-s1">jax</span>.<span class="pl-s1">config</span>.<span class="pl-en">update</span>(<span class="pl-s">'jax_numpy_rank_promotion'</span>, <span class="pl-s">'raise'</span>)
<span class="pl-k">from</span> <span class="pl-s1">jax</span> <span class="pl-k">import</span> <span class="pl-s1">lax</span>, <span class="pl-s1">vmap</span>

<span class="pl-k">import</span> <span class="pl-s1">math</span>
<span class="pl-k">import</span> <span class="pl-s1">numpy</span> <span class="pl-k">as</span> <span class="pl-s1">onp</span>

<span class="pl-k">import</span> <span class="pl-s1">einops</span>


<span class="pl-c"># # Grid Packing / Unpacking</span>

<span class="pl-k">def</span> <span class="pl-en">get_grid_centers</span>(<span class="pl-s1">box_size</span>, <span class="pl-s1">cell_size</span>, <span class="pl-s1">num_dims</span>):
  <span class="pl-s1">cells_count_per_side</span> <span class="pl-c1">=</span> <span class="pl-en">int</span>(<span class="pl-s1">onp</span>.<span class="pl-en">ceil</span>(<span class="pl-s1">box_size</span> <span class="pl-c1">/</span> <span class="pl-s1">cell_size</span>))

  <span class="pl-s1">grid_centers</span> <span class="pl-c1">=</span> <span class="pl-s1">onp</span>.<span class="pl-en">zeros</span>(<span class="pl-en">tuple</span>(
      [<span class="pl-s1">cells_count_per_side</span>]<span class="pl-c1">*</span><span class="pl-s1">num_dims</span> <span class="pl-c1">+</span> [<span class="pl-s1">num_dims</span>]
  ))

  <span class="pl-k">for</span> <span class="pl-s1">i</span> <span class="pl-c1">in</span> <span class="pl-en">range</span>(<span class="pl-s1">num_dims</span>):
    <span class="pl-s1">grid_centers</span>[...,<span class="pl-s1">i</span>] <span class="pl-c1">=</span> (<span class="pl-s1">onp</span>.<span class="pl-en">reshape</span>(<span class="pl-s1">onp</span>.<span class="pl-en">arange</span>(<span class="pl-s1">cells_count_per_side</span>),
                                       <span class="pl-en">tuple</span>([<span class="pl-c1">1</span>]<span class="pl-c1">*</span><span class="pl-s1">i</span> <span class="pl-c1">+</span> [<span class="pl-s1">cells_count_per_side</span>] <span class="pl-c1">+</span> [<span class="pl-c1">1</span>]<span class="pl-c1">*</span>(<span class="pl-s1">num_dims</span><span class="pl-c1">-</span><span class="pl-s1">i</span><span class="pl-c1">-</span><span class="pl-c1">1</span>)))
                                       <span class="pl-c1">+</span> <span class="pl-c1">0.5</span>)<span class="pl-c1">*</span><span class="pl-s1">cell_size</span>

  <span class="pl-k">return</span> <span class="pl-s1">np</span>.<span class="pl-en">array</span>(<span class="pl-s1">grid_centers</span>)


<span class="pl-k">def</span> <span class="pl-en">positions_to_grid</span>(<span class="pl-s1">position</span>, <span class="pl-s1">box_size</span>, <span class="pl-s1">cell_size</span>, <span class="pl-s1">num_dims</span>, <span class="pl-s1">particle_id</span>, <span class="pl-s1">aux</span><span class="pl-c1">=</span><span class="pl-c1">None</span>):
  <span class="pl-s1">cells</span> <span class="pl-c1">=</span> <span class="pl-en">int</span>(<span class="pl-s1">math</span>.<span class="pl-en">ceil</span>(<span class="pl-s1">box_size</span> <span class="pl-c1">/</span> <span class="pl-s1">cell_size</span>))
  <span class="pl-v">N</span> <span class="pl-c1">=</span> <span class="pl-s1">position</span>.<span class="pl-s1">shape</span>[<span class="pl-c1">0</span>]

  <span class="pl-s1">cell_contents</span> <span class="pl-c1">=</span> <span class="pl-s1">np</span>.<span class="pl-en">concatenate</span>([<span class="pl-s1">position</span>, <span class="pl-s1">particle_id</span>[:, <span class="pl-c1">None</span>]], <span class="pl-s1">axis</span><span class="pl-c1">=</span><span class="pl-c1">-</span><span class="pl-c1">1</span>)
  
  <span class="pl-k">if</span> <span class="pl-s1">aux</span> <span class="pl-c1">is</span> <span class="pl-c1">not</span> <span class="pl-c1">None</span>:
    <span class="pl-s1">cell_contents</span>, <span class="pl-c1">*</span><span class="pl-s1">_</span> <span class="pl-c1">=</span> <span class="pl-en">set_aux</span>(<span class="pl-s1">cell_contents</span>, <span class="pl-s1">aux</span>)

  <span class="pl-s1">grid</span> <span class="pl-c1">=</span> <span class="pl-s1">np</span>.<span class="pl-en">zeros</span>((<span class="pl-s1">cells</span> <span class="pl-c1">**</span> <span class="pl-s1">num_dims</span>, <span class="pl-s1">cell_contents</span>.<span class="pl-s1">shape</span>[<span class="pl-c1">-</span><span class="pl-c1">1</span>]))
  <span class="pl-s1">grid</span> <span class="pl-c1">=</span> <span class="pl-s1">grid</span>.<span class="pl-s1">at</span>[:<span class="pl-v">N</span>, :].<span class="pl-en">set</span>(<span class="pl-s1">cell_contents</span>)

  <span class="pl-s1">centers</span> <span class="pl-c1">=</span> <span class="pl-en">get_grid_centers</span>(<span class="pl-s1">box_size</span>, <span class="pl-s1">cell_size</span>, <span class="pl-s1">num_dims</span>).<span class="pl-en">reshape</span>((<span class="pl-c1">-</span><span class="pl-c1">1</span>, <span class="pl-s1">num_dims</span>))
  <span class="pl-s1">grid</span> <span class="pl-c1">=</span> <span class="pl-s1">grid</span>.<span class="pl-s1">at</span>[:<span class="pl-v">N</span>, :<span class="pl-s1">num_dims</span>].<span class="pl-en">add</span>(<span class="pl-c1">-</span><span class="pl-s1">centers</span>[:<span class="pl-v">N</span>])

  <span class="pl-s1">grid</span> <span class="pl-c1">=</span> <span class="pl-s1">np</span>.<span class="pl-en">reshape</span>(<span class="pl-s1">grid</span>, (<span class="pl-s1">cells</span>,) <span class="pl-c1">*</span> <span class="pl-s1">num_dims</span> <span class="pl-c1">+</span>  (<span class="pl-c1">-</span><span class="pl-c1">1</span>,))

  <span class="pl-k">return</span> <span class="pl-s1">grid</span>


<span class="pl-k">def</span> <span class="pl-en">pad_axis_channel_last</span>(<span class="pl-v">X</span>, <span class="pl-s1">factors</span>, <span class="pl-s1">axis_ind</span>, <span class="pl-s1">max_grid_distance</span> <span class="pl-c1">=</span> <span class="pl-c1">1</span>):

  <span class="pl-s1">num_dims</span> <span class="pl-c1">=</span> <span class="pl-en">len</span>(<span class="pl-s1">factors</span>)
  <span class="pl-s1">axes_s</span> <span class="pl-c1">=</span> <span class="pl-s">' '</span>.<span class="pl-en">join</span>([<span class="pl-en">chr</span>(<span class="pl-en">ord</span>(<span class="pl-s">'x'</span>) <span class="pl-c1">+</span> <span class="pl-s1">i</span>) <span class="pl-k">for</span> <span class="pl-s1">i</span> <span class="pl-c1">in</span> <span class="pl-en">range</span>(<span class="pl-s1">num_dims</span>)])
  <span class="pl-s1">factors_s</span> <span class="pl-c1">=</span> <span class="pl-s">' '</span>.<span class="pl-en">join</span>([<span class="pl-s">'f'</span><span class="pl-c1">+</span><span class="pl-en">chr</span>(<span class="pl-en">ord</span>(<span class="pl-s">'x'</span>) <span class="pl-c1">+</span> <span class="pl-s1">i</span>) <span class="pl-k">for</span> <span class="pl-s1">i</span> <span class="pl-c1">in</span> <span class="pl-en">range</span>(<span class="pl-s1">num_dims</span>)])
  <span class="pl-s1">fact_d</span> <span class="pl-c1">=</span> {}
  <span class="pl-k">for</span> <span class="pl-s1">i</span> <span class="pl-c1">in</span> <span class="pl-en">range</span>(<span class="pl-s1">num_dims</span>):
    <span class="pl-s1">fact_d</span>[<span class="pl-s">'f'</span> <span class="pl-c1">+</span> <span class="pl-en">chr</span>(<span class="pl-en">ord</span>(<span class="pl-s">'x'</span>) <span class="pl-c1">+</span> <span class="pl-s1">i</span>)] <span class="pl-c1">=</span> <span class="pl-s1">factors</span>[<span class="pl-s1">i</span>]

  <span class="pl-c"># all entries for all dimensions before the target axis</span>
  <span class="pl-s1">all_slice</span> <span class="pl-c1">=</span> <span class="pl-en">tuple</span>(<span class="pl-en">slice</span>(<span class="pl-v">X</span>.<span class="pl-s1">shape</span>[<span class="pl-s1">j</span>]) <span class="pl-k">for</span> <span class="pl-s1">j</span> <span class="pl-c1">in</span> <span class="pl-en">range</span>(<span class="pl-s1">axis_ind</span> <span class="pl-c1">+</span> <span class="pl-c1">1</span>))

  <span class="pl-c"># the end halo along the target axis</span>
  <span class="pl-s1">idx</span> <span class="pl-c1">=</span> <span class="pl-s1">all_slice</span> <span class="pl-c1">+</span> (<span class="pl-en">slice</span>(<span class="pl-c1">-</span><span class="pl-s1">max_grid_distance</span>, <span class="pl-v">X</span>.<span class="pl-s1">shape</span>[<span class="pl-s1">axis_ind</span> <span class="pl-c1">+</span> <span class="pl-c1">1</span>]),)
  <span class="pl-s1">p</span> <span class="pl-c1">=</span> <span class="pl-v">X</span>[<span class="pl-s1">idx</span>]
  <span class="pl-s1">p</span> <span class="pl-c1">=</span> <span class="pl-s1">einops</span>.<span class="pl-en">rearrange</span>(<span class="pl-s1">p</span>, <span class="pl-s">f'(<span class="pl-s1"><span class="pl-kos">{</span><span class="pl-s1">factors_s</span><span class="pl-kos">}</span></span>) <span class="pl-s1"><span class="pl-kos">{</span><span class="pl-s1">axes_s</span><span class="pl-kos">}</span></span> c -&gt; <span class="pl-s1"><span class="pl-kos">{</span><span class="pl-s1">factors_s</span><span class="pl-kos">}</span></span> <span class="pl-s1"><span class="pl-kos">{</span><span class="pl-s1">axes_s</span><span class="pl-kos">}</span></span> c'</span>, <span class="pl-c1">**</span><span class="pl-s1">fact_d</span>)
  <span class="pl-s1">p</span> <span class="pl-c1">=</span> <span class="pl-s1">np</span>.<span class="pl-en">roll</span>(<span class="pl-s1">p</span>, (<span class="pl-c1">0</span>,) <span class="pl-c1">*</span> <span class="pl-s1">axis_ind</span> <span class="pl-c1">+</span> (<span class="pl-c1">1</span>,) <span class="pl-c1">+</span> (<span class="pl-c1">0</span>,) <span class="pl-c1">*</span> (<span class="pl-v">X</span>.<span class="pl-s1">ndim</span> <span class="pl-c1">-</span> <span class="pl-s1">axis_ind</span> <span class="pl-c1">-</span> <span class="pl-c1">1</span>), <span class="pl-s1">axis</span><span class="pl-c1">=</span><span class="pl-en">range</span>(<span class="pl-v">X</span>.<span class="pl-s1">ndim</span>))
  <span class="pl-s1">p1</span> <span class="pl-c1">=</span> <span class="pl-s1">einops</span>.<span class="pl-en">rearrange</span>(<span class="pl-s1">p</span>, <span class="pl-s">f'<span class="pl-s1"><span class="pl-kos">{</span><span class="pl-s1">factors_s</span><span class="pl-kos">}</span></span> <span class="pl-s1"><span class="pl-kos">{</span><span class="pl-s1">axes_s</span><span class="pl-kos">}</span></span> c -&gt; (<span class="pl-s1"><span class="pl-kos">{</span><span class="pl-s1">factors_s</span><span class="pl-kos">}</span></span>) <span class="pl-s1"><span class="pl-kos">{</span><span class="pl-s1">axes_s</span><span class="pl-kos">}</span></span> c'</span>)

  <span class="pl-c"># the start halo along the target axis</span>
  <span class="pl-s1">idx</span> <span class="pl-c1">=</span> <span class="pl-s1">all_slice</span> <span class="pl-c1">+</span> (<span class="pl-en">slice</span>(<span class="pl-c1">0</span>, <span class="pl-s1">max_grid_distance</span>),)
  <span class="pl-s1">p</span> <span class="pl-c1">=</span> <span class="pl-v">X</span>[<span class="pl-s1">idx</span>]
  <span class="pl-s1">p</span> <span class="pl-c1">=</span> <span class="pl-s1">einops</span>.<span class="pl-en">rearrange</span>(<span class="pl-s1">p</span>, <span class="pl-s">f'(<span class="pl-s1"><span class="pl-kos">{</span><span class="pl-s1">factors_s</span><span class="pl-kos">}</span></span>) <span class="pl-s1"><span class="pl-kos">{</span><span class="pl-s1">axes_s</span><span class="pl-kos">}</span></span> c -&gt; <span class="pl-s1"><span class="pl-kos">{</span><span class="pl-s1">factors_s</span><span class="pl-kos">}</span></span> <span class="pl-s1"><span class="pl-kos">{</span><span class="pl-s1">axes_s</span><span class="pl-kos">}</span></span> c'</span>, <span class="pl-c1">**</span><span class="pl-s1">fact_d</span>)
  <span class="pl-s1">p</span> <span class="pl-c1">=</span> <span class="pl-s1">np</span>.<span class="pl-en">roll</span>(<span class="pl-s1">p</span>, (<span class="pl-c1">0</span>,) <span class="pl-c1">*</span> <span class="pl-s1">axis_ind</span> <span class="pl-c1">+</span> (<span class="pl-c1">-</span><span class="pl-c1">1</span>,) <span class="pl-c1">+</span> (<span class="pl-c1">0</span>,) <span class="pl-c1">*</span> (<span class="pl-v">X</span>.<span class="pl-s1">ndim</span> <span class="pl-c1">-</span> <span class="pl-s1">axis_ind</span> <span class="pl-c1">-</span> <span class="pl-c1">1</span>), <span class="pl-s1">axis</span><span class="pl-c1">=</span><span class="pl-en">range</span>(<span class="pl-v">X</span>.<span class="pl-s1">ndim</span>))
  <span class="pl-s1">p2</span> <span class="pl-c1">=</span> <span class="pl-s1">einops</span>.<span class="pl-en">rearrange</span>(<span class="pl-s1">p</span>, <span class="pl-s">f'<span class="pl-s1"><span class="pl-kos">{</span><span class="pl-s1">factors_s</span><span class="pl-kos">}</span></span> <span class="pl-s1"><span class="pl-kos">{</span><span class="pl-s1">axes_s</span><span class="pl-kos">}</span></span> c -&gt; (<span class="pl-s1"><span class="pl-kos">{</span><span class="pl-s1">factors_s</span><span class="pl-kos">}</span></span>) <span class="pl-s1"><span class="pl-kos">{</span><span class="pl-s1">axes_s</span><span class="pl-kos">}</span></span> c'</span>)

  <span class="pl-v">X</span> <span class="pl-c1">=</span> <span class="pl-s1">np</span>.<span class="pl-en">concatenate</span>([<span class="pl-s1">p1</span>, <span class="pl-v">X</span>, <span class="pl-s1">p2</span>], <span class="pl-s1">axis</span><span class="pl-c1">=</span><span class="pl-c1">1</span> <span class="pl-c1">+</span> <span class="pl-s1">axis_ind</span>)

  <span class="pl-k">return</span> <span class="pl-v">X</span>

<span class="pl-k">def</span> <span class="pl-en">fold_factors</span>(<span class="pl-s1">batch_size</span>, <span class="pl-s1">grid_shape</span>, <span class="pl-s1">max_grid_distance</span>):
  <span class="pl-s1">num_dims</span> <span class="pl-c1">=</span> <span class="pl-en">len</span>(<span class="pl-s1">grid_shape</span>)
  <span class="pl-s1">max_folds</span> <span class="pl-c1">=</span> <span class="pl-s1">onp</span>.<span class="pl-en">log</span>(<span class="pl-s1">batch_size</span>) <span class="pl-c1">/</span> <span class="pl-s1">onp</span>.<span class="pl-en">log</span>(<span class="pl-c1">2</span>)

  <span class="pl-s1">folds</span> <span class="pl-c1">=</span> <span class="pl-s1">onp</span>.<span class="pl-en">zeros</span>((<span class="pl-s1">num_dims</span>,), <span class="pl-s1">onp</span>.<span class="pl-s1">int32</span>)

  <span class="pl-k">for</span> <span class="pl-s1">i</span> <span class="pl-c1">in</span> <span class="pl-en">range</span>(<span class="pl-en">int</span>(<span class="pl-s1">max_folds</span>)<span class="pl-c1">*</span><span class="pl-s1">num_dims</span>):
    <span class="pl-s1">dim</span> <span class="pl-c1">=</span> <span class="pl-s1">i</span> <span class="pl-c1">%</span> <span class="pl-s1">num_dims</span>
    <span class="pl-c"># folds[dim]+2, rather than +1, since we need an even number of cells for pairwise exchange in update_grid_locations</span>
    <span class="pl-s1">width</span>, <span class="pl-s1">remain</span> <span class="pl-c1">=</span> <span class="pl-en">divmod</span>(<span class="pl-s1">grid_shape</span>[<span class="pl-s1">dim</span>], <span class="pl-c1">2</span><span class="pl-c1">**</span>(<span class="pl-s1">folds</span>[<span class="pl-s1">dim</span>]<span class="pl-c1">+</span><span class="pl-c1">2</span>))
    <span class="pl-c"># width, remain = divmod(grid_shape[dim], 2**(folds[dim]+1))</span>
    <span class="pl-c"># TODO is 2*max_grid_distance the right cutoff?</span>
    <span class="pl-k">if</span> <span class="pl-s1">remain</span> <span class="pl-c1">==</span> <span class="pl-c1">0</span> <span class="pl-c1">and</span> <span class="pl-s1">width</span> <span class="pl-c1">&gt;=</span> <span class="pl-s1">max_grid_distance</span>:
      <span class="pl-s1">folds</span>[<span class="pl-s1">dim</span>] <span class="pl-c1">+=</span> <span class="pl-c1">1</span>
    <span class="pl-k">if</span> <span class="pl-s1">onp</span>.<span class="pl-en">sum</span>(<span class="pl-s1">folds</span>) <span class="pl-c1">==</span> <span class="pl-s1">max_folds</span>:
      <span class="pl-k">break</span>
  <span class="pl-s1">factors</span> <span class="pl-c1">=</span> <span class="pl-c1">2</span> <span class="pl-c1">**</span> <span class="pl-s1">folds</span>

  <span class="pl-k">if</span> <span class="pl-s1">onp</span>.<span class="pl-en">sum</span>(<span class="pl-s1">folds</span>) <span class="pl-c1">&lt;</span> <span class="pl-s1">max_folds</span>:
    <span class="pl-en">print</span>(<span class="pl-s">f"Folds <span class="pl-s1"><span class="pl-kos">{</span><span class="pl-s1">folds</span><span class="pl-kos">}</span></span> have sum smaller than target of <span class="pl-s1"><span class="pl-kos">{</span><span class="pl-s1">max_folds</span><span class="pl-kos">}</span></span>. This corresponds to a batch size of <span class="pl-s1"><span class="pl-kos">{</span><span class="pl-s1">onp</span>.<span class="pl-en">prod</span>(<span class="pl-s1">factors</span>)<span class="pl-kos">}</span></span> rather than the target of <span class="pl-s1"><span class="pl-kos">{</span><span class="pl-s1">batch_size</span><span class="pl-kos">}</span></span>."</span>)

  <span class="pl-k">return</span> <span class="pl-s1">factors</span>

<span class="pl-k">def</span> <span class="pl-en">order_grid_by_factors</span>(<span class="pl-s1">num_dims</span>):
  <span class="pl-s1">front_order</span> <span class="pl-c1">=</span> ()
  <span class="pl-s1">back_order</span> <span class="pl-c1">=</span> ()

  <span class="pl-k">for</span> <span class="pl-s1">i</span> <span class="pl-c1">in</span> <span class="pl-en">range</span>(<span class="pl-s1">num_dims</span>):
    <span class="pl-s1">front_order</span> <span class="pl-c1">+=</span> (<span class="pl-c1">2</span> <span class="pl-c1">*</span> <span class="pl-s1">i</span>,)
    <span class="pl-s1">back_order</span> <span class="pl-c1">+=</span> (<span class="pl-c1">2</span> <span class="pl-c1">*</span> <span class="pl-s1">i</span> <span class="pl-c1">+</span> <span class="pl-c1">1</span>,)

  <span class="pl-k">return</span> <span class="pl-s1">front_order</span> <span class="pl-c1">+</span> <span class="pl-s1">back_order</span> <span class="pl-c1">+</span> (<span class="pl-c1">2</span> <span class="pl-c1">*</span> <span class="pl-s1">num_dims</span>,)

<span class="pl-k">def</span> <span class="pl-en">fold_grid</span>(<span class="pl-v">X</span>, <span class="pl-s1">max_grid_distance</span>, <span class="pl-s1">batch_size</span><span class="pl-c1">=</span><span class="pl-c1">None</span>, <span class="pl-s1">factors</span><span class="pl-c1">=</span><span class="pl-c1">None</span>, <span class="pl-s1">inner_fold</span><span class="pl-c1">=</span><span class="pl-c1">True</span>):
  <span class="pl-s1">num_dims</span> <span class="pl-c1">=</span> <span class="pl-v">X</span>.<span class="pl-s1">ndim</span> <span class="pl-c1">-</span> <span class="pl-c1">1</span>

  <span class="pl-k">if</span> <span class="pl-s1">factors</span> <span class="pl-c1">is</span> <span class="pl-c1">None</span>:
    <span class="pl-s1">factors</span> <span class="pl-c1">=</span> <span class="pl-en">fold_factors</span>(<span class="pl-s1">batch_size</span>, <span class="pl-v">X</span>.<span class="pl-s1">shape</span>[:<span class="pl-c1">-</span><span class="pl-c1">1</span>], <span class="pl-s1">max_grid_distance</span>)

  <span class="pl-v">X_shape</span> <span class="pl-c1">=</span> []
  <span class="pl-k">for</span> <span class="pl-s1">i</span> <span class="pl-c1">in</span> <span class="pl-en">range</span>(<span class="pl-s1">num_dims</span>):
    <span class="pl-s1">_</span>, <span class="pl-s1">ragged</span> <span class="pl-c1">=</span> <span class="pl-en">divmod</span>(<span class="pl-v">X</span>.<span class="pl-s1">shape</span>[<span class="pl-s1">i</span>], <span class="pl-s1">factors</span>[<span class="pl-s1">i</span>])
    <span class="pl-k">assert</span> <span class="pl-s1">ragged</span> <span class="pl-c1">==</span> <span class="pl-c1">0</span>
    <span class="pl-v">X_shape</span> <span class="pl-c1">+=</span> [<span class="pl-s1">factors</span>[<span class="pl-s1">i</span>], <span class="pl-v">X</span>.<span class="pl-s1">shape</span>[<span class="pl-s1">i</span>] <span class="pl-c1">//</span> <span class="pl-s1">factors</span>[<span class="pl-s1">i</span>]]
  <span class="pl-v">X_shape</span> <span class="pl-c1">+=</span> [<span class="pl-v">X</span>.<span class="pl-s1">shape</span>[<span class="pl-c1">-</span><span class="pl-c1">1</span>]]

  <span class="pl-k">if</span> <span class="pl-s1">onp</span>.<span class="pl-en">min</span>(<span class="pl-s1">onp</span>.<span class="pl-en">array</span>(<span class="pl-v">X</span>.<span class="pl-s1">shape</span>[:<span class="pl-c1">-</span><span class="pl-c1">1</span>]) <span class="pl-c1">/</span> <span class="pl-s1">onp</span>.<span class="pl-en">array</span>(<span class="pl-s1">factors</span>)) <span class="pl-c1">&lt;</span> <span class="pl-s1">max_grid_distance</span>:
    <span class="pl-en">print</span>(<span class="pl-s">f"Folded grid may be too small. This failure is happening when splitting across devices. factors <span class="pl-s1"><span class="pl-kos">{</span><span class="pl-s1">factors</span><span class="pl-kos">}</span></span> X.shape <span class="pl-s1"><span class="pl-kos">{</span><span class="pl-v">X</span>.<span class="pl-s1">shape</span><span class="pl-kos">}</span></span> new X_shape <span class="pl-s1"><span class="pl-kos">{</span><span class="pl-v">X_shape</span><span class="pl-kos">}</span></span>."</span>)

  <span class="pl-s1">ordering</span> <span class="pl-c1">=</span> <span class="pl-en">order_grid_by_factors</span>(<span class="pl-s1">num_dims</span>)

  <span class="pl-v">X</span> <span class="pl-c1">=</span> <span class="pl-s1">np</span>.<span class="pl-en">reshape</span>(<span class="pl-v">X</span>, <span class="pl-v">X_shape</span>)
  <span class="pl-v">X</span> <span class="pl-c1">=</span> <span class="pl-s1">np</span>.<span class="pl-en">transpose</span>(<span class="pl-v">X</span>, <span class="pl-s1">ordering</span>)

  <span class="pl-k">if</span> <span class="pl-s1">inner_fold</span>:
    <span class="pl-v">X</span> <span class="pl-c1">=</span> <span class="pl-v">X</span>.<span class="pl-en">reshape</span>((<span class="pl-c1">-</span><span class="pl-c1">1</span>,) <span class="pl-c1">+</span> <span class="pl-v">X</span>.<span class="pl-s1">shape</span>[<span class="pl-s1">num_dims</span>:])
  <span class="pl-k">else</span>:
    <span class="pl-v">X</span> <span class="pl-c1">=</span> <span class="pl-en">folded_pad</span>(<span class="pl-v">X</span>, <span class="pl-s1">max_grid_distance</span>, <span class="pl-s1">factors</span>)

  <span class="pl-k">return</span> <span class="pl-v">X</span>, <span class="pl-s1">factors</span>

<span class="pl-c"># Dynamics</span>

<span class="pl-k">def</span> <span class="pl-en">pairwise_exchange</span>(<span class="pl-v">X</span>, <span class="pl-s1">cell_size</span>, <span class="pl-s1">axis</span>, <span class="pl-s1">num_dims</span>):
  <span class="pl-s1">particle_data_size</span> <span class="pl-c1">=</span> <span class="pl-v">X</span>.<span class="pl-s1">shape</span>[<span class="pl-c1">-</span><span class="pl-c1">1</span>] <span class="pl-c1">-</span> <span class="pl-s1">num_dims</span>

  <span class="pl-s1">axes_pre</span> <span class="pl-c1">=</span> <span class="pl-s">' '</span>.<span class="pl-en">join</span>([<span class="pl-s">'('</span> <span class="pl-c1">+</span> <span class="pl-en">chr</span>(<span class="pl-en">ord</span>(<span class="pl-s">'x'</span>) <span class="pl-c1">+</span> <span class="pl-s1">i</span>) <span class="pl-c1">+</span> <span class="pl-s">' q)'</span> <span class="pl-k">if</span> <span class="pl-s1">i</span> <span class="pl-c1">==</span> <span class="pl-s1">axis</span> <span class="pl-k">else</span> <span class="pl-en">chr</span>(<span class="pl-en">ord</span>(<span class="pl-s">'x'</span>) <span class="pl-c1">+</span> <span class="pl-s1">i</span>) <span class="pl-k">for</span> <span class="pl-s1">i</span> <span class="pl-c1">in</span> <span class="pl-en">range</span>(<span class="pl-s1">num_dims</span>)])
  <span class="pl-s1">axes_post</span> <span class="pl-c1">=</span> <span class="pl-s">' '</span>.<span class="pl-en">join</span>([<span class="pl-en">chr</span>(<span class="pl-en">ord</span>(<span class="pl-s">'x'</span>) <span class="pl-c1">+</span> <span class="pl-s1">i</span>) <span class="pl-k">for</span> <span class="pl-s1">i</span> <span class="pl-c1">in</span> <span class="pl-en">range</span>(<span class="pl-s1">num_dims</span>)])
  <span class="pl-v">X</span> <span class="pl-c1">=</span> <span class="pl-s1">einops</span>.<span class="pl-en">rearrange</span>(<span class="pl-v">X</span>, <span class="pl-s">f'<span class="pl-s1"><span class="pl-kos">{</span><span class="pl-s1">axes_pre</span><span class="pl-kos">}</span></span> c -&gt; q <span class="pl-s1"><span class="pl-kos">{</span><span class="pl-s1">axes_post</span><span class="pl-kos">}</span></span> c'</span>, <span class="pl-s1">q</span><span class="pl-c1">=</span><span class="pl-c1">2</span>)

  <span class="pl-s1">shift</span> <span class="pl-c1">=</span> <span class="pl-s1">onp</span>.<span class="pl-en">zeros</span>((<span class="pl-c1">2</span>, <span class="pl-s1">num_dims</span><span class="pl-c1">+</span><span class="pl-s1">particle_data_size</span>))
  <span class="pl-s1">shift</span>[<span class="pl-c1">0</span>,<span class="pl-s1">axis</span>] <span class="pl-c1">=</span> <span class="pl-s1">cell_size</span>
  <span class="pl-s1">shift</span>[<span class="pl-c1">1</span>,<span class="pl-s1">axis</span>] <span class="pl-c1">=</span> <span class="pl-c1">-</span><span class="pl-s1">cell_size</span>
  <span class="pl-s1">shift</span> <span class="pl-c1">=</span> <span class="pl-s1">shift</span>.<span class="pl-en">reshape</span>((<span class="pl-c1">2</span>,) <span class="pl-c1">+</span> (<span class="pl-c1">1</span>,)<span class="pl-c1">*</span><span class="pl-s1">num_dims</span> <span class="pl-c1">+</span> (<span class="pl-s1">num_dims</span><span class="pl-c1">+</span><span class="pl-s1">particle_data_size</span>,))
  <span class="pl-s1">_Xrev</span> <span class="pl-c1">=</span> <span class="pl-v">X</span>[::<span class="pl-c1">-</span><span class="pl-c1">1</span>]
  <span class="pl-v">X_swapped</span> <span class="pl-c1">=</span> <span class="pl-s1">_Xrev</span> <span class="pl-c1">+</span> <span class="pl-s1">shift</span><span class="pl-c1">*</span>(<span class="pl-s1">_Xrev</span>[..., [<span class="pl-c1">-</span><span class="pl-c1">1</span>]]<span class="pl-c1">&gt;</span><span class="pl-c1">0</span>) <span class="pl-c"># only add offset to occupied cells</span>

  <span class="pl-s1">square_displacement</span> <span class="pl-c1">=</span> <span class="pl-k">lambda</span> <span class="pl-s1">_X</span>: <span class="pl-s1">np</span>.<span class="pl-en">sum</span>(<span class="pl-s1">np</span>.<span class="pl-en">sum</span>(<span class="pl-s1">_X</span>[...,:<span class="pl-s1">num_dims</span>]<span class="pl-c1">**</span><span class="pl-c1">2</span> <span class="pl-c1">*</span> (<span class="pl-s1">_X</span>[...,[<span class="pl-c1">-</span><span class="pl-c1">1</span>]]<span class="pl-c1">&gt;</span><span class="pl-c1">0</span>), <span class="pl-s1">axis</span><span class="pl-c1">=</span><span class="pl-c1">-</span><span class="pl-c1">1</span>, <span class="pl-s1">keepdims</span><span class="pl-c1">=</span><span class="pl-c1">True</span>), <span class="pl-s1">axis</span><span class="pl-c1">=</span><span class="pl-c1">0</span>, <span class="pl-s1">keepdims</span><span class="pl-c1">=</span><span class="pl-c1">True</span>)

  <span class="pl-s1">keep_mask</span> <span class="pl-c1">=</span> (<span class="pl-en">square_displacement</span>(<span class="pl-v">X</span>) <span class="pl-c1">&lt;</span> <span class="pl-en">square_displacement</span>(<span class="pl-v">X_swapped</span>))

  <span class="pl-v">X_new</span> <span class="pl-c1">=</span> <span class="pl-v">X</span><span class="pl-c1">*</span><span class="pl-s1">keep_mask</span> <span class="pl-c1">+</span> <span class="pl-v">X_swapped</span><span class="pl-c1">*</span>(<span class="pl-c1">~</span><span class="pl-s1">keep_mask</span>)

  <span class="pl-v">X_new</span> <span class="pl-c1">=</span> <span class="pl-s1">einops</span>.<span class="pl-en">rearrange</span>(<span class="pl-v">X_new</span>, <span class="pl-s">f'q <span class="pl-s1"><span class="pl-kos">{</span><span class="pl-s1">axes_post</span><span class="pl-kos">}</span></span> c -&gt; <span class="pl-s1"><span class="pl-kos">{</span><span class="pl-s1">axes_pre</span><span class="pl-kos">}</span></span> c'</span>)

  <span class="pl-k">return</span> <span class="pl-v">X_new</span>

<span class="pl-s1">pairwise_exchange</span> <span class="pl-c1">=</span> <span class="pl-en">vmap</span>(<span class="pl-s1">pairwise_exchange</span>, <span class="pl-s1">in_axes</span><span class="pl-c1">=</span>(<span class="pl-c1">0</span>, <span class="pl-c1">None</span>, <span class="pl-c1">None</span>, <span class="pl-c1">None</span>))


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

  <span class="pl-s1">axis</span> <span class="pl-c1">=</span> <span class="pl-c1">1</span>
  <span class="pl-v">X</span> <span class="pl-c1">=</span> <span class="pl-en">pad_axis_channel_last</span>(<span class="pl-v">X</span>, <span class="pl-s1">factors</span>, <span class="pl-s1">axis</span>)
  <span class="pl-v">X</span> <span class="pl-c1">=</span> <span class="pl-en">pairwise_exchange</span>(<span class="pl-v">X</span>, <span class="pl-s1">cell_size</span>, <span class="pl-s1">axis</span>, <span class="pl-s1">num_dims</span>)
  <span class="pl-s1">body_slice</span> <span class="pl-c1">=</span> <span class="pl-en">tuple</span>(<span class="pl-en">slice</span>(<span class="pl-v">X</span>.<span class="pl-s1">shape</span>[<span class="pl-s1">j</span>]) <span class="pl-k">for</span> <span class="pl-s1">j</span> <span class="pl-c1">in</span> <span class="pl-en">range</span>(<span class="pl-s1">axis</span><span class="pl-c1">+</span><span class="pl-c1">1</span>)) <span class="pl-c1">+</span> (<span class="pl-en">slice</span>(<span class="pl-c1">1</span>,<span class="pl-c1">-</span><span class="pl-c1">1</span>),)
  <span class="pl-v">X</span> <span class="pl-c1">=</span> <span class="pl-v">X</span>[<span class="pl-s1">body_slice</span>]

  <span class="pl-k">return</span> <span class="pl-v">X</span>, <span class="pl-c1">None</span>

<span class="pl-k">def</span> <span class="pl-en">test_grid</span>(<span class="pl-s1">grid_topology</span><span class="pl-c1">=</span>(), <span class="pl-s1">num_dims</span><span class="pl-c1">=</span><span class="pl-c1">2</span>):

  <span class="pl-s1">cell_size</span> <span class="pl-c1">=</span> <span class="pl-c1">1.</span><span class="pl-c1">/</span><span class="pl-c1">4.</span>
  <span class="pl-s1">max_grid_distance</span> <span class="pl-c1">=</span> <span class="pl-c1">5</span>
  <span class="pl-s1">box_size</span> <span class="pl-c1">=</span> <span class="pl-c1">512.</span>

  <span class="pl-s1">points</span> <span class="pl-c1">=</span> <span class="pl-s1">onp</span>.<span class="pl-en">linspace</span>(<span class="pl-c1">0.</span>, <span class="pl-c1">50</span>, <span class="pl-s1">num</span><span class="pl-c1">=</span><span class="pl-c1">66</span>)
  <span class="pl-v">X</span>, <span class="pl-v">Y</span> <span class="pl-c1">=</span> <span class="pl-s1">onp</span>.<span class="pl-en">meshgrid</span>(<span class="pl-s1">points</span>, <span class="pl-s1">points</span>)
  <span class="pl-v">R</span> <span class="pl-c1">=</span> <span class="pl-s1">onp</span>.<span class="pl-en">stack</span>((<span class="pl-v">X</span>.<span class="pl-en">ravel</span>(), <span class="pl-v">Y</span>.<span class="pl-en">ravel</span>()), <span class="pl-s1">axis</span><span class="pl-c1">=</span><span class="pl-c1">1</span>) <span class="pl-c1">+</span> <span class="pl-c1">0.1</span>

  <span class="pl-v">R</span> <span class="pl-c1">=</span> <span class="pl-s1">onp</span>.<span class="pl-en">array</span>(<span class="pl-v">R</span>, <span class="pl-s1">onp</span>.<span class="pl-s1">float64</span>)

  <span class="pl-s1">grid</span> <span class="pl-c1">=</span> <span class="pl-en">positions_to_grid</span>(<span class="pl-v">R</span>, <span class="pl-s1">box_size</span>, <span class="pl-s1">cell_size</span>, <span class="pl-s1">num_dims</span>, <span class="pl-s1">np</span>.<span class="pl-en">arange</span>(<span class="pl-v">R</span>.<span class="pl-s1">shape</span>[<span class="pl-c1">0</span>]) <span class="pl-c1">+</span> <span class="pl-c1">1</span>)

  <span class="pl-s1">grid</span>, <span class="pl-s1">factors</span> <span class="pl-c1">=</span> <span class="pl-en">fold_grid</span>(<span class="pl-s1">grid</span>, <span class="pl-s1">max_grid_distance</span>,
                                        <span class="pl-s1">batch_size</span><span class="pl-c1">=</span><span class="pl-c1">128</span>)

  <span class="pl-k">def</span> <span class="pl-en">cond_fn</span>(<span class="pl-s1">grid_old_grid</span>):
    <span class="pl-s1">grid</span>, <span class="pl-s1">i</span> <span class="pl-c1">=</span> <span class="pl-s1">grid_old_grid</span>
    <span class="pl-k">return</span> (<span class="pl-s1">i</span> <span class="pl-c1">&lt;</span> <span class="pl-c1">1</span>)

  <span class="pl-k">def</span> <span class="pl-en">body_fn</span>(<span class="pl-s1">grid_old_grid</span>):
    <span class="pl-s1">grid</span>, <span class="pl-s1">i</span> <span class="pl-c1">=</span> <span class="pl-s1">grid_old_grid</span>
    <span class="pl-s1">_grid</span>, <span class="pl-s1">overflow_count</span> <span class="pl-c1">=</span> <span class="pl-en">update_grid_locations</span>(<span class="pl-s1">grid</span>, <span class="pl-s1">cell_size</span>, <span class="pl-s1">num_dims</span>, <span class="pl-s1">factors</span>)
    <span class="pl-k">return</span> <span class="pl-s1">_grid</span>, <span class="pl-s1">i</span><span class="pl-c1">+</span><span class="pl-c1">1</span>

  <span class="pl-s1">grid_old_grid</span> <span class="pl-c1">=</span> (<span class="pl-s1">grid</span>, <span class="pl-c1">0</span>)

  <span class="pl-c">### BUG IS HERE. Compute a while loop two different ways. Result should be the</span>
  <span class="pl-c">### same. Instead it's different.</span>

  <span class="pl-s1">new_tuple_lax</span> <span class="pl-c1">=</span> <span class="pl-s1">lax</span>.<span class="pl-en">while_loop</span>(<span class="pl-s1">cond_fn</span>, <span class="pl-s1">body_fn</span>, <span class="pl-s1">grid_old_grid</span>)

  <span class="pl-s1">new_tuple_while</span> <span class="pl-c1">=</span> <span class="pl-s1">grid_old_grid</span>
  <span class="pl-k">while</span> <span class="pl-en">cond_fn</span>(<span class="pl-s1">new_tuple_while</span>):
    <span class="pl-s1">new_tuple_while</span> <span class="pl-c1">=</span> <span class="pl-en">body_fn</span>(<span class="pl-s1">new_tuple_while</span>)

  <span class="pl-s1">max_dev</span> <span class="pl-c1">=</span> <span class="pl-s1">onp</span>.<span class="pl-en">max</span>(<span class="pl-s1">onp</span>.<span class="pl-en">abs</span>(<span class="pl-s1">new_tuple_lax</span>[<span class="pl-c1">0</span>] <span class="pl-c1">-</span> <span class="pl-s1">new_tuple_while</span>[<span class="pl-c1">0</span>]))
  <span class="pl-en">print</span>(<span class="pl-s1">max_dev</span>)
  <span class="pl-k">assert</span> <span class="pl-s1">max_dev</span> <span class="pl-c1">==</span> <span class="pl-c1">0</span>

<span class="pl-en">test_grid</span>()</pre></div>
<p dir="auto">Hi all,</p>
<p dir="auto">I was discussing with <a class="user-mention notranslate" data-hovercard-type="user" data-hovercard-url="/users/jakevdp/hovercard" data-octo-click="hovercard-link-click" data-octo-dimensions="link_type:self" href="https://github.com/jakevdp">@jakevdp</a> in <a class="issue-link js-issue-link" data-error-text="Failed to load title" data-id="859207960" data-permission-text="Title is private" data-url="https://github.com/google/jax/issues/6466" data-hovercard-type="pull_request" data-hovercard-url="/google/jax/pull/6466/hovercard" href="https://github.com/google/jax/pull/6466">#6466</a> 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:</p>
<blockquote>
<ul dir="auto">
<li>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.</li>
<li>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</li>
<li>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</li>
<li>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.</li>
</ul>
</blockquote>
<p dir="auto">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?</p>
<p dir="auto">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 <a href="https://colab.research.google.com/drive/1dyfhD4SaSAfKPhx-aG7lF8BY88WlsWWx?usp=sharing" rel="nofollow">notebook</a> 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!</p>
<table role="table">
<thead>
<tr>
<th align="left">Feature</th>
<th>Pytorch</th>
<th>Tensorflow</th>
</tr>
</thead>
<tbody>
<tr>
<td align="left">Documentation</td>
<td><a href="https://pytorch.org/docs/stable/sparse.html#" rel="nofollow">here</a></td>
<td><a href="https://www.tensorflow.org/guide/sparse_tensor" rel="nofollow">here</a></td>
</tr>
<tr>
<td align="left">Formats</td>
<td>COO</td>
<td>COO</td>
</tr>
<tr>
<td align="left">Supported ops</td>
<td>variations of matmul</td>
<td>sums, min, matmul, element-wise ops through <code class="notranslate">tf.sparse.map_values</code></td>
</tr>
<tr>
<td align="left">Specialized ops</td>
<td>Probably<sup><a href="#user-content-fn-5-dc92f6f795e789591e336bd2748f0620" id="user-content-fnref-5-dc92f6f795e789591e336bd2748f0620" data-footnote-ref="" aria-describedby="footnote-label">1</a></sup></td>
<td>Seems so <sup><a href="#user-content-fn-4-dc92f6f795e789591e336bd2748f0620" id="user-content-fnref-4-dc92f6f795e789591e336bd2748f0620" data-footnote-ref="" aria-describedby="footnote-label">2</a></sup></td>
</tr>
<tr>
<td align="left">Grad of sparse<sup><a href="#user-content-fn-1-dc92f6f795e789591e336bd2748f0620" id="user-content-fnref-1-dc92f6f795e789591e336bd2748f0620" data-footnote-ref="" aria-describedby="footnote-label">3</a></sup></td>
<td>Limited to specific matmul</td>
<td>No.</td>
</tr>
<tr>
<td align="left">Sparse grad<sup><a href="#user-content-fn-2-dc92f6f795e789591e336bd2748f0620" id="user-content-fnref-2-dc92f6f795e789591e336bd2748f0620" data-footnote-ref="" aria-describedby="footnote-label">4</a></sup></td>
<td>Yes</td>
<td>No.</td>
</tr>
<tr>
<td align="left">Dimensions</td>
<td>2D, Hybrid</td>
<td>2D - tensor supports higher-D, but operations not</td>
</tr>
<tr>
<td align="left">Uncoalesced<sup><a href="#user-content-fn-3-dc92f6f795e789591e336bd2748f0620" id="user-content-fnref-3-dc92f6f795e789591e336bd2748f0620" data-footnote-ref="" aria-describedby="footnote-label">5</a></sup>allowed</td>
<td>Yes</td>
<td>No mention</td>
</tr>
<tr>
<td align="left">Sparse - sparse matmul</td>
<td>No</td>
<td>No</td>
</tr>
<tr>
<td align="left">Extra goodies</td>
<td>Has an adam optimizer for sparse tensor</td>
<td>Nope</td>
</tr>
</tbody>
</table>
<p dir="auto">To do</p>
<ul dir="auto">
<li>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.)</li>
<li>Check how pydata-sparse implement n-D sparse arrays.</li>
</ul>
<p dir="auto">Proposal:</p>
<ul dir="auto">
<li>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)</li>
</ul>
<p dir="auto">I figured we can use this issue to track progress and discuss, so ... thoughts?</p>
<section data-footnotes="" class="footnotes"><h2 id="footnote-label" class="sr-only" dir="auto">Footnotes</h2>
<ol dir="auto">
<li id="user-content-fn-5-dc92f6f795e789591e336bd2748f0620">
<p dir="auto">They mention Torchcudasparse is now defunct <a href="https://github.com/pytorch/pytorch/tree/master/aten/src">here</a> but I can't find any other mention in the source code. <a href="#user-content-fnref-5-dc92f6f795e789591e336bd2748f0620" data-footnote-backref="" aria-label="Back to reference 1" class="data-footnote-backref"><g-emoji class="g-emoji" alias="leftwards_arrow_with_hook" fallback-src="https://github.githubassets.com/images/icons/emoji/unicode/21a9.png">↩</g-emoji></a></p>
</li>
<li id="user-content-fn-4-dc92f6f795e789591e336bd2748f0620">
<p dir="auto">Sparse ops are imported from gen_sparse_ops.py, which doesnt exist in the repo? <a href="#user-content-fnref-4-dc92f6f795e789591e336bd2748f0620" data-footnote-backref="" aria-label="Back to reference 2" class="data-footnote-backref"><g-emoji class="g-emoji" alias="leftwards_arrow_with_hook" fallback-src="https://github.githubassets.com/images/icons/emoji/unicode/21a9.png">↩</g-emoji></a></p>
</li>
<li id="user-content-fn-1-dc92f6f795e789591e336bd2748f0620">
<p dir="auto">Meaning gradient w.r.t to the sparse tensor <a href="#user-content-fnref-1-dc92f6f795e789591e336bd2748f0620" data-footnote-backref="" aria-label="Back to reference 3" class="data-footnote-backref"><g-emoji class="g-emoji" alias="leftwards_arrow_with_hook" fallback-src="https://github.githubassets.com/images/icons/emoji/unicode/21a9.png">↩</g-emoji></a></p>
</li>
<li id="user-content-fn-2-dc92f6f795e789591e336bd2748f0620">
<p dir="auto">Meaning calculating the gradient without turning it into a dense tensor <a href="#user-content-fnref-2-dc92f6f795e789591e336bd2748f0620" data-footnote-backref="" aria-label="Back to reference 4" class="data-footnote-backref"><g-emoji class="g-emoji" alias="leftwards_arrow_with_hook" fallback-src="https://github.githubassets.com/images/icons/emoji/unicode/21a9.png">↩</g-emoji></a></p>
</li>
<li id="user-content-fn-3-dc92f6f795e789591e336bd2748f0620">
<p dir="auto">Uncoalesced means duplicate entries in the coordinates - total value of the element is then sum of duplicate elements. <a href="#user-content-fnref-3-dc92f6f795e789591e336bd2748f0620" data-footnote-backref="" aria-label="Back to reference 5" class="data-footnote-backref"><g-emoji class="g-emoji" alias="leftwards_arrow_with_hook" fallback-src="https://github.githubassets.com/images/icons/emoji/unicode/21a9.png">↩</g-emoji></a></p>
</li>
</ol>
</section>