[mahdi-shafiei]

Summary

 smiles_to_graph(smiles)
   mol = RDKit parse
-  node_features = jnp.stack([jnp.array(atom features) for each atom])
-  for each bond: four .at[i, j].set(...) on the (N, N) and (N, N, 4) arrays
+  host arrays filled with NumPy (same feature lists, same bond one-hots)
+  one jnp.asarray per array
 batch_smiles_to_graphs(smiles_list)
-  jnp.zeros + .at[i, :n].set(...) per molecule   (float64 under x64)
+  NumPy padding, one jnp.asarray per array      (float32 always)

smiles_to_graph built every graph with eager JAX ops on host data: about 105 device dispatches per BBBP molecule, each bond's scatter copying the whole adjacency or edge-feature array. The graph is now assembled with NumPy and converted once per array. get_atom_features and get_bond_features keep their signatures and share the host-side feature lists with the graph builders.

Evidence

Same output. Over all 1,640 BBBP SMILES, the NumPy path gives the same dtype, shape, values and num_nodes as the shipped code for the 1,629 valid molecules, and rejects the same 11.

Tests. tests/operators/drug_discovery/test_molecular_graphs.py, written first:

• Before: 3 failed. Under x64 a batch was float64. An empty batch failed inside max(), and an atomless SMILES failed inside jnp.stack.
• After: 8 passed. Exact ethanol and benzene graphs, batch rows equal to the padded single graphs, float32 JAX arrays with and without x64, and refusals that name their cause.

The guard tests passed before and after the change. The 113 drug-discovery tests (tests/operators/drug_discovery, tests/operators/test_drug_discovery.py) and the MolNet and DTI benchmark tests (15) pass, and pre-commit run --all-files passes all 24 hooks.

Speed. MolNetBenchmark(quick=True).run(), the fixture that sets the CI benchmarks shard's floor, timed in one fresh CPU process per run, alternating main and this branch: main 59.3 s and 60.8 s, this branch 13.0 s and 13.4 s. In CI's last two benchmark-shard runs that fixture took 110.1 s and 249.2 s.

Merge Danger

Door: two-way.

Blast Radius: drug-discovery graph inputs.

The graphs are unchanged for every molecule checked. The behaviour changes are: a batch under x64 is float32 instead of float64 (DiffBio's tests run without x64), and an empty batch or an atomless SMILES raises a ValueError naming its cause instead of one from max() or jnp.stack. No repository outside DiffBio calls these functions (all 3,916 workspace repositories checked).