how to make Structure object available during model inference

[breadbread1984]

I found the class Structure object are saved in class BatchDict featurised_example at https://github.com/google-deepmind/alphafold3/blob/main/run_alphafold.py#L437 . but function utils.remove_invalidly_typed_feats remove it because its value is an object rather than jnp.ndarray.

I tried removing utils.remove_invalidly_typed_feats or add back cleaned_struc back to featurised_example but both causes JAX errors.

I want to use bonds table in the Structure object, how can I pass it into model?

[amyssnippet]

Hi @breadbread1984,

This error occurs because JAX functions (specifically when JIT-compiled) require inputs to be valid "Pytrees" (nested dictionaries, lists, or tuples of arrays). They cannot handle arbitrary Python objects (like the Structure class) unless those classes are explicitly registered as Pytree nodes.

The function remove_invalidly_typed_feats is there specifically to sanitize the input and prevent these non-array objects from causing JAX tracing errors downstream.

Recommended Strategy:
Instead of passing the entire Structure object, you should extract the specific data you need (the bonds table) and convert it into a standard array.

1. Extract & Convert: Before the cleanup function runs, extract the bonds table from the Structure object and convert it into a numpy.ndarray (e.g., an array of shape [N_bonds, 2]).
2. Add to Dictionary: Add this new array to the featurised_example dictionary with a new key (e.g., 'bond_indices').
3. Pass it: Since it is now a standard array, remove_invalidly_typed_feats will not remove it, and the JAX model will be able to accept it as a valid input feature.

Note: If your model code expects fixed-size inputs (due to jit), ensure that your bonds array is padded to a consistent size if the number of bonds varies between examples.