Quick fix to reallocation with nd.array-type box

[arturtoshev]

Minimal working example:

import jax.numpy as jnp
from jax_md import space, partition

pos = jnp.array([[0.0, 0.0]])
box_size = jnp.array([1.0,  1.0])
cell_size = 0.1

displacement_fn, shift_fn = space.periodic(box_size)
neighbor_list_fn = partition.neighbor_list(
    displacement_fn, box_size, cell_size)

print('First iteration neighbor allocation')
_ = neighbor_list_fn.allocate(pos)
print('Second iteration neighbor allocation')
_ = neighbor_list_fn.allocate(pos)

leads to

First iteration neighbor allocation
Second iteration neighbor allocation
Traceback (most recent call last):
  File "/home/artur/code/jaxmdQF/test.py", line 15, in <module>
    _ = neighbor_list_fn.allocate(pos)
  File "/home/artur/code/jaxmdQF/jax-md/jax_md/partition.py", line 981, in allocate_fn
    return neighbor_list_fn(position, extra_capacity=extra_capacity, **kwargs)
  File "/home/artur/code/jaxmdQF/jax-md/jax_md/partition.py", line 953, in neighbor_list_fn
    return neighbor_fn((position, PartitionError(jnp.zeros((), jnp.uint8))))
  File "/home/artur/code/jaxmdQF/jax-md/jax_md/partition.py", line 910, in neighbor_fn
    idx = cell_list_candidate_fn(cl, position)
ValueError: The truth value of an array with more than one element is ambiguous. Use a.any() or a.all()

I'm not really sure why this happens, but my quick fix is the only way I could get around the error message. I backtracked the error and it originates from this commit: 511eb19

Note: this bug happens only when box_size is an array (both np.ndarray and jnp.ndarray).

[arturtoshev]

One can also safely use an older jax-md version. The latest version before the above-mentioned commit is 0.2.3.

PS: Of course, I mean ndarray in the title.