Migratimg MPNN tutorial to Keras3

[maitry63]

This PR migrates the Message Passing Neural Network (MPNN) molecular property prediction tutorial from Keras 2 to Keras 3 and built on top of PR #2289.

Original contribution by @LakshmiKalaKadali.

[gemini-code-assist]

Code Review

This pull request updates the Message-Passing Neural Network (MPNN) molecular property prediction example to Keras 3, making it backend-agnostic (supporting JAX, PyTorch, and TensorFlow) by replacing tf.RaggedTensor with fixed-size padded buffers and masks, implementing a custom keras.utils.PyDataset, and using a gated readout with hybrid pooling. The code review highlighted several critical issues in both the script and notebook versions: a compilation crash under TensorFlow due to comparing symbolic shapes with Python's max function, potential runtime crashes from unhandled None values when parsing invalid SMILES strings (both during featurization and test visualization), biased training metrics caused by zero-padding the final batch in the dataset, and an incorrect learning rate schedule calculation that uses batch size instead of steps per epoch.

[maitry63]

@gemini-code-assist review

[gemini-code-assist]

Code Review

This pull request migrates the MPNN molecular graphs example to Keras 3, making it backend-agnostic to support JAX, PyTorch, and TensorFlow. It replaces TensorFlow-specific ragged tensors with fixed-size buffers, introduces a pre-featurization step to eliminate CPU bottlenecks, and implements a more stable gated readout. The review feedback highlights several critical issues: a bug in the dataset indexing that shifts padding bonds and corrupts atom features, a masking flaw in the gated readout where segment max pooling fails to handle negative features correctly, the inclusion of failed molecules as noisy training placeholders, and redundant code in the bond featurizer that can be simplified.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

[gemini-code-assist]

In MPNNDataset.__getitem__, shifting all bond indices with batch_pair_indices[i] += i * MAX_ATOMS also shifts the unused/padding bonds (which are initialized to [0, 0]). As a result, these padding bonds are mapped to [i * MAX_ATOMS, i * MAX_ATOMS], which represents valid self-loops on the first atom of each molecule.

During message passing, because EdgeNetwork includes a bias term, these padding bonds will produce non-zero messages (equal to bias * neighbor_feat). These messages are then summed into the first atom's features via ops.segment_sum. This means the first atom of each molecule will receive a large number of spurious messages proportional to the number of padding bonds, corrupting the graph representation and making it dependent on the padding size.

To fix this, consider introducing a bond_mask in the dataset and passing it to the model to mask out the messages of padding bonds before they are summed.

[gemini-code-assist]

In GatedReadout.call, gated_x is multiplied by mask to zero out the features of padded atoms. However, because tanh (used in self.feat) can output negative values, the valid features in gated_x can be negative (e.g., -0.5). Since padded atoms have their features set to 0.0, ops.segment_max will incorrectly select 0.0 (from the padded atoms) instead of the correct negative maximum from the valid atoms (since 0.0 > -0.5).

To prevent padded elements from interfering with the max pooling, we should mask the padded elements with a large negative value (e.g., -1e9) instead of 0.0 before calling ops.segment_max.

Suggested change

	x_max = ops.segment_max(
	gated_x,
	ops.cast(indicator, "int32"),
	num_segments=num_molecules,
	)
	x_max = ops.segment_max(
	gated_x + (1.0 - mask) * -1e9,
	ops.cast(indicator, "int32"),
	num_segments=num_molecules,
	)

[gemini-code-assist]

In the pre-featurization loop, if smiles_to_graph returns None (because RDKit failed to parse the SMILES), a placeholder tuple of all-zero arrays is appended to processed_data.

While the atom mask will zero out the features during readout, the model is still trained on these all-zero graphs with their real labels (self.labels[real_idx]). This introduces noise into the training process and can degrade model performance.

Instead of keeping placeholders for failed molecules, we should filter them out of the dataset entirely during the pre-featurization step and update the dataframe df accordingly.

# Pre-featurize once to remove the RDKit bottleneck during training.
print("Pre-featurizing Dataset...")
processed_data = []
valid_indices = []
for i, smiles_string in enumerate(tqdm(df.smiles.values)):
    graph = smiles_to_graph(smiles_string)
    if graph is not None:
        processed_data.append(graph)
        valid_indices.append(i)

df = df.iloc[valid_indices].reset_index(drop=True)

[gemini-code-assist]

The loop in BondFeaturizer.encode is identical to the one in Featurizer.encode. Since self.dim in BondFeaturizer is already updated to include the extra dimension for None bonds, we can simplify this by calling super().encode(bond) to avoid code duplication.

        return super().encode(bond)