XlaRuntimeError: PyTree custom to_iterable function should return a pair

[TomClarkMassSpec]

Hi,

I’m having trouble escaping the same error after several revisions of the code, using a custom pytree. I have a pytree with four nodes. The first node is what my predict functions are differentiated with respect to. The remaining three nodes are auxiliary which I’ve now made explicit, yet the same error relentlessly surfaces with respect to the third node whenever I try to create an array of pytrees from an array of the third nodes (protein_arguments_tuple, which is a tuple of integers all of equal length). It does work with a single pytree made with a single protein_arguments_tuple as the third node. I abandoned the array in favour of a list of pytrees and a for loop which works but I am just kicking the can down the road because when I try to map my predict function to the list of pytrees, I get the exact same error. The error appears to be with the flatten_customPyTree ?

Error: “…. File "/opt/anaconda3/lib/python3.8/site-packages/jax/_src/tree_util.py", line 58, in tree_flatten
return pytree.flatten(tree, is_leaf)
XlaRuntimeError: PyTree custom to_iterable function should return a pair”

I also tried giving it ‘a pair’ in the form of a Tuple, None and Tuple, False to no avail.

I seem to get things to work with slow for loops. The idea here is with no grad in use, and the variables held constant (a single epoch), the predict function is to be applied to the number of data points; hence, the list of protein_arguments_tuple which I couldn’t make into in either a JAX array or numpy array or put into an earlier single function without creating arrays or lists.

Relevant current code snippets:
def mark_auxiliary_data(x):
    return x, None

class CustomPyTree:
    def __init__(self, variables_List, constants_tuple, protein_arguments_tuple, sds_normalization_factor_list):
        self.variables_List = variables_List
        self.constants_tuple = tree_util.tree_map(mark_auxiliary_data, constants_tuple)
        self.protein_arguments_tuple = tree_util.tree_map(mark_auxiliary_data, protein_arguments_tuple)
        self.sds_normalization_factor_list = tree_util.tree_map(mark_auxiliary_data, sds_normalization_factor_list)

def flatten_CustomPyTree(container) -> Tuple[Tuple[float], Tuple[float, int], Tuple[int], Tuple[float]]:
    return container.variables_List, container.constants_tuple, container.protein_arguments_tuple, container.sds_normalization_factor_list

def unflatten_CustomPyTree(tree_def, leaves) -> CustomPyTree:
    variables_List, constants_tuple, protein_arguments_tuple, sds_normalization_factor_list = leaves
    return CustomPyTree(variables_List, constants_tuple, protein_arguments_tuple, sds_normalization_factor_list)

# register custom pytree using register_pytree_node from jax.tree_util
tree_util.register_pytree_node(
    CustomPyTree,
    flatten_CustomPyTree,
    unflatten_CustomPyTree
)

…..

def build_custom_pytree(variables_List, constants_tuple, protein_arguments_tuple, sds_normalization_factor_list):
    custom_tree = CustomPyTree(variables_List, constants_tuple, protein_arguments_tuple, sds_normalization_factor_list)
    return custom_tree

def build_a_list_of_pytrees(variables_List, constants_tuple, x_seq_array, sds_normalization_factor_list): #tip arrays dont work easily
    list_of_pytrees = []
    for protein_arguments_tuple in x_seq_array:
        list_of_pytrees.append(build_custom_pytree(variables_List, constants_tuple, protein_arguments_tuple, sds_normalization_factor_list))
    #array_of_pytrees = jnp.asarray(list_of_pytrees)
    return list_of_pytrees

list_of_pytrees = build_a_list_of_pytrees(variables_List, constants_tuple, x_seq_array, sds_normalization_factor_list)

def apply_function_to_pytrees(function, list_of_pytrees):
    return [jax.tree_util.tree_map(function, pytree) for pytree in list_of_pytrees]

processed_pytrees_for_pending_loss_function = apply_function_to_pytrees(distance_prediction, list_of_pytrees)

Thanks for looking,
Tom

[mattjj]

Thanks for the question! But IIUC this is a copy of #16119, so maybe we should keep the discussion on that thread. To do that, I'll close this thread.