Numpy dtype changes pointer when pickling causing pointer comparison to fail in `compute_atom_dim`

In chemprop/chemprop#1386, we found that cuik-molmaker featurization doesn't work with multiple dataloader workers on MacOS. The symptom is the resulting batch of atom features has a shape n_atoms x 0. (Same for bond features.)

I believe this issue is this:

• On Mac, pytorch's DataLoader uses spawn to handle the new processes. This uses pickle to copy the dataset to each process. (link)
• Our datasets contain the featurizer, so the cuik_molmaker featurizer is also pickled
• The cuik_molmaker featurizer contains numpy arrays produced by e.g., atom_onehot_feature_names_to_array. These arrays have a dtype of numpy.int64.
• Usually there is only one numpy.int64 in memory and all dtype check compare pointers to that one instance in memory (singleton). But pickle needs to own the data it is using so makes another copy of numpy.int64 in memory.
• So after a pickle/unpickle roundtrip, the numpy array will appear to have the original dtype, but the pointer is different.
• In compute_atom_dim , it checks atom_property_list_onehot.dtype() == py::dtype("int64") https://github.com/NVIDIA-BioNeMo/cuik-molmaker/blob/a81100a1f240a0bd0a372c2e817aa0d80c7dca51/src/features.cpp#L475-L478 which compares pointers. The pointers don't match so the atom feature dim is set to 0, hench our problem.
• Note: The operator == in pybind11, is equivalent to python's is pointer comparison. (link)

So my question is if we need to have the dtype check? If so, maybe we should have it give an error if it fails the dtype check instead of silently returning 0. To fix the dtype check, an LLM suggested this:
Change
atom_property_list_onehot.dtype() == py::dtype("int64")
to

py::detail::npy_api::get().PyArray_EquivTypes_(
    atom_property_list_onehot.dtype().ptr(), 
    py::dtype("int64").ptr()
)

(I don't understand c++ well enough myself to know if this is a real solution.)
Or maybe we leave off the dtype check entirely. This applies for both one hot and float properties in both compute_atom_dim and compute_bond_dim.

Here is some code I used to test things:

import cuik_molmaker
import pickle
A = cuik_molmaker.atom_onehot_feature_names_to_array(["atomic-number-common"])
atom_feats, *_ = cuik_molmaker.batch_mol_featurizer(["C", "CC"], A, [], [], False, False, False, False)
print(atom_feats.shape)
A_pkl = pickle.loads(pickle.dumps(A))
atom_feats, *_ = cuik_molmaker.batch_mol_featurizer(["C", "CC"], A_pkl, [], [], False, False, False, False)
print(atom_feats.shape)
atom_feats, *_ = cuik_molmaker.batch_mol_featurizer(["C", "CC"], A, [], [], False, False, False, False)
print(atom_feats.shape)
print(A, A_pkl)
print(type(A), type(A_pkl))
print(type(A[0]), type(A_pkl[0]))

gives

(3, 38)
(3, 0)
(3, 38)
[1] [1]
<class 'numpy.ndarray'> <class 'numpy.ndarray'>
<class 'numpy.int64'> <class 'numpy.int64'>

d = pickle.loads(pickle.dumps(np.dtype('int64')))
print(d == np.dtype('int64'))
print(d is np.dtype('int64'))

gives

True
False

From an LLM:

import pickle, numpy as np, cuik_molmaker

A = cuik_molmaker.atom_onehot_feature_names_to_array(["atomic-number-common"])

def dim(a):
    f, *_ = cuik_molmaker.batch_mol_featurizer(["C", "CC"], a, [], [], False, False, False, False)
    return f.shape[1]

def probe(name, a):
    print(f"{name:12s} dtype={a.dtype!s:>7} num={a.dtype.num} bo={a.dtype.byteorder} "
          f"shape={a.shape} strides={a.strides} C={a.flags['C_CONTIGUOUS']} "
          f"W={a.flags['WRITEABLE']} O={a.flags['OWNDATA']} AL={a.flags['ALIGNED']} "
          f"align8={a.ctypes.data % 8} base={type(a.base).__name__} dim={dim(a)}")

# 1) Is it pickling, or just "the 2nd call"?
print("repeat calls with A:", dim(A), dim(A), dim(A))

# 2) Protocol dependence  (numpy uses PickleBuffer/np.frombuffer for protocol >= 5)
for p in (2, 4, 5):
    probe(f"proto{p}", pickle.loads(pickle.dumps(A, protocol=p)))

# 3) What "fixes" it?
B = pickle.loads(pickle.dumps(A))
probe("pkl", B)
probe("pkl.copy", B.copy())
probe("ascontig", np.ascontiguousarray(B))
probe("astype", B.astype(np.int64))
probe("handmade", np.array([1], dtype=np.int64))
probe("orig", A)

gives

repeat calls with A: 38 38 38
proto2       dtype=  int64 num=7 bo== shape=(1,) strides=(8,) C=True W=True O=True AL=True align8=0 base=NoneType dim=0
proto4       dtype=  int64 num=7 bo== shape=(1,) strides=(8,) C=True W=True O=True AL=True align8=0 base=NoneType dim=0
proto5       dtype=  int64 num=7 bo== shape=(1,) strides=(8,) C=True W=True O=False AL=True align8=0 base=ndarray dim=0
pkl          dtype=  int64 num=7 bo== shape=(1,) strides=(8,) C=True W=True O=True AL=True align8=0 base=NoneType dim=0
pkl.copy     dtype=  int64 num=7 bo== shape=(1,) strides=(8,) C=True W=True O=True AL=True align8=0 base=NoneType dim=0
ascontig     dtype=  int64 num=7 bo== shape=(1,) strides=(8,) C=True W=True O=True AL=True align8=0 base=NoneType dim=0
astype       dtype=  int64 num=7 bo== shape=(1,) strides=(8,) C=True W=True O=True AL=True align8=0 base=NoneType dim=38
handmade     dtype=  int64 num=7 bo== shape=(1,) strides=(8,) C=True W=True O=True AL=True align8=0 base=NoneType dim=38
orig         dtype=  int64 num=7 bo== shape=(1,) strides=(8,) C=True W=True O=False AL=True align8=0 base=PyCapsule dim=38

The relevant entry points are compute_atom_dim and compute_bond_dim in src/features.cpp; first reproduce the pickle round-trip with the Python snippet in the issue and inspect the dtype comparisons. Done means both atom and bond dimensions remain correct after unpickling, with an explicit outcome for unsupported dtypes rather than a silent zero.