NL precomputing doesn't work for periodic systems of different size.

[M-R-Schaefer]

I was just trying to fit the HM21 dataset and it appears that neighborlists are not correctly reallocated for differently sized periodic system (works for gas phase)

│ 975 in neighbor_list_fn                                                                          │
│                                                                                                  │
│    972 │   d = partial(metric_sq, **kwargs)                                                      │
│    973 │   d = vmap(d)                                                                           │
│    974 │   return lax.cond(                                                                      │
│ ❱  975 │   │   jnp.any(d(position, nbrs.reference_position) > threshold_sq),                     │
│    976 │   │   (position, nbrs.error), neighbor_fn,                                              │
│    977 │   │   nbrs, lambda x: x)                                                                │
│    978                                                                                           │
╰──────────────────────────────────────────────────────────────────────────────────────────────────╯
ValueError: vmap got inconsistent sizes for array axes to be mapped:
  * one axis had size 8: axis 0 of argument Ra of type float64[8,3];
  * one axis had size 32: axis 0 of argument Rb of type float64[32,3]