Neighborlist shape changes when updated with 2D box

[pl-fuchs]

When using a box with box.ndim == 2 to update the neighborlist, the error-leaf of the neighbor list changes from uint8[] to uint8[3,3].
I think this is due to a check of the box format returning a matrix instead of a scalar in partition.py:

def is_box_valid(box: Array) -> bool:
  if jnp.isscalar(box) or box.ndim == 0 or box.ndim == 1:
    return True
  if box.ndim == 2:
    return jnp.triu(box) == box
  return False

at

https://github.com/jax-md/jax-md/blob/4861632f1478d9ba3207d0a8bc308969afc1486a/jax_md/partition.py#L587C6-L587C6

Is it possible that there might be missing a jnp.all(...)? I patched this function in my code and now the neighborlist does not change its shape anymore.